Make BloqBuilder.join accept array-like of soquets

quantumlib · Dec 18, 2024 · 37e9f7b · 37e9f7b
1 parent 47a36c8
commit 37e9f7b
Show file tree

Hide file tree

Showing 2 changed files with 9 additions and 2 deletions.
diff --git a/qualtran/_infra/composite_bloq.py b/qualtran/_infra/composite_bloq.py
@@ -1234,12 +1234,13 @@ def split(self, soq: Soquet) -> NDArray[Soquet]:  # type: ignore[type-var]
 
         return self.add(Split(dtype=soq.reg.dtype), reg=soq)
 
-    def join(self, soqs: NDArray[Soquet], dtype: Optional[QDType] = None) -> Soquet:  # type: ignore[type-var]
+    def join(self, soqs: SoquetInT, dtype: Optional[QDType] = None) -> Soquet:
         from qualtran.bloqs.bookkeeping import Join
 
         try:
+            soqs = np.asarray(soqs)
             (n,) = soqs.shape
-        except AttributeError:
+        except (AttributeError, ValueError):
             raise ValueError("`join` expects a 1-d array of input soquets to join.") from None
 
         if not all(soq.reg.bitsize == 1 for soq in soqs):

diff --git a/qualtran/_infra/composite_bloq_test.py b/qualtran/_infra/composite_bloq_test.py
@@ -389,6 +389,12 @@ def test_util_convenience_methods():
     assert len(cbloq.connections) == 1 + 10 + 1
 
 
+def test_join_list():
+    bb = BloqBuilder()
+    qs = [bb.allocate() for _ in range(10)]
+    _ = bb.join(qs)
+
+
 def test_util_convenience_methods_errors():
     bb = BloqBuilder()