Changes to .flatten

 - Default predicate to flatten everything
 - If you can't flatten a bloq, don't complain; just don't do anything
 - Documentation improvements

Author: mpharrigan

qualtran/_infra/composite_bloq.ipynb

@@ -256,8 +256,8 @@
     "cbloq2 = cbloq.copy()\n",
     "\n",
     "# They're the same!\n",
-    "display(show_bloq(cbloq))\n",
-    "display(show_bloq(cbloq2))"
+    "show_bloq(cbloq)\n",
+    "show_bloq(cbloq2)"
    ]
   },
   {
@@ -364,7 +364,7 @@
     "# right-dangling soquets.\n",
     "fsoqs = bb.map_soqs(cbloq.final_soqs(), soq_map)\n",
     "copy = bb.finalize(**fsoqs)\n",
-    "copy"
+    "print(copy)"
    ]
   },
   {
@@ -475,6 +475,8 @@
     "        stuff = bb.add(TestParallelCombo(), reg=stuff)\n",
     "        return {'stuff': stuff}\n",
     "\n",
+    "# Note! We're using `.as_composite_bloq()` to wrap the Bloq\n",
+    "# into a compute graph with one node.\n",
     "three_p = ThreeParallelBloqs().as_composite_bloq()\n",
     "show_bloq(three_p)"
    ]
@@ -486,8 +488,9 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "# Do one flattening iteration\n",
-    "flat_three_p = three_p.flatten_once(lambda binst: True)\n",
+    "# Do one flattening operation; here equivalent to `.decompose_bloq()`\n",
+    "# on the original bloq.\n",
+    "flat_three_p = three_p.flatten_once()\n",
     "show_bloq(flat_three_p)"
    ]
   },
@@ -496,7 +499,7 @@
    "id": "6f322361",
    "metadata": {},
    "source": [
-    "If we just decomposed all of the subbloqs, our diagram would look very similar except the subbloqs are now `CompositeBloq` container classes instead of the original `TestParallelBloq`."
+    "Now we have a `CompositeBloq` with three subbloqs. What if we wanted to continue decomposing? A naive approach would be to simply call `.decompose_bloq()` on each subbloqs. We'll do this in the next cell and see what happens."
    ]
   },
   {
@@ -523,6 +526,14 @@
     "show_bloq(decompose_children)"
    ]
   },
+  {
+   "cell_type": "markdown",
+   "id": "35b8a7d3-d54e-4b72-a50a-4fb047bf797c",
+   "metadata": {},
+   "source": [
+    "This is still a `CompositeBloq` with three subbloqs. The only difference is now the subbloqs are each compute graphs of their own. This likely isn't what we want. Instead, we want to do the equivalent of `flatMap` with the decompose operation: namely, decompose the subbloqs and remove a layer of nesting. This is what the `flatten_once` method achieves."
+   ]
+  },
   {
    "cell_type": "code",
    "execution_count": null,
@@ -532,10 +543,18 @@
    "source": [
     "# Actually do a flattening operation on all the subbloqs\n",
     "show_bloq(\n",
-    "    flat_three_p.flatten_once(lambda binst: True)\n",
+    "    flat_three_p.flatten_once()\n",
     ")"
    ]
   },
+  {
+   "cell_type": "markdown",
+   "id": "34e0e027-f327-47a4-9169-de233cdb8359",
+   "metadata": {},
+   "source": [
+    "You can use the optional predicate to control which subbloqs get decomposed and flattened."
+   ]
+  },
   {
    "cell_type": "code",
    "execution_count": null,
@@ -549,16 +568,25 @@
     ")"
    ]
   },
+  {
+   "cell_type": "markdown",
+   "id": "f13ff34a-a9f3-43e6-85bd-4a1f90f4e3e1",
+   "metadata": {},
+   "source": [
+    "The `.flatten` method will repeatedly call `flatten_once` until you can't flatten any more."
+   ]
+  },
   {
    "cell_type": "code",
    "execution_count": null,
    "id": "c1e6cbb0",
    "metadata": {},
    "outputs": [],
    "source": [
-    "# Flatten until you can't flatten any more\n",
+    "# Note that in this example, we have gone back to the original `three_p` starting composite bloq.\n",
+    "# This will perform two flattening operations.\n",
     "show_bloq(\n",
-    "    three_p.flatten(lambda binst: binst.bloq.supports_decompose_bloq())\n",
+    "    three_p.flatten()\n",
     ")"
    ]
   }
@@ -579,7 +607,7 @@
    "name": "python",
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
-   "version": "3.10.9"
+   "version": "3.11.8"
   }
  },
  "nbformat": 4,

qualtran/_infra/composite_bloq.py

@@ -40,7 +40,7 @@
 import sympy
 from numpy.typing import NDArray
 
-from .bloq import Bloq, DecomposeTypeError
+from .bloq import Bloq, DecomposeNotImplementedError, DecomposeTypeError
 from .data_types import check_dtypes_consistent, QAny, QBit, QDType
 from .quantum_graph import BloqInstance, Connection, DanglingT, LeftDangle, RightDangle, Soquet
 from .registers import Register, Side, Signature
@@ -316,7 +316,9 @@ def copy(self) -> 'CompositeBloq':
         fsoqs = _map_soqs(self.final_soqs(), soq_map)
         return bb.finalize(**fsoqs)
 
-    def flatten_once(self, pred: Callable[[BloqInstance], bool]) -> 'CompositeBloq':
+    def flatten_once(
+        self, pred: Callable[[BloqInstance], bool] = lambda binst: True
+    ) -> 'CompositeBloq':
         """Decompose and flatten each subbloq that satisfies `pred`.
 
         This will only flatten "once". That is, we will go through the bloq instances
@@ -326,16 +328,17 @@ def flatten_once(self, pred: Callable[[BloqInstance], bool]) -> 'CompositeBloq':
         Args:
             pred: A predicate that takes a bloq instance and returns True if it should
                 be decomposed and flattened or False if it should remain undecomposed.
-                All bloqs for which this callable returns True must support decomposition.
+                If the bloq does not have a decomposition, it will remain undecomposed.
+                By default, flatten everything.
 
         Returns:
             A new composite bloq where subbloqs matching `pred` have been decomposed and
             flattened.
 
         Raises:
-            NotImplementedError: If `pred` returns True but the underlying bloq does not
-                support `decompose_bloq()`.
-            DidNotFlattenAnythingError: If none of the bloq instances satisfied `pred`.
+            DidNotFlattenAnythingError: If the operation did not actually flatten anything.
+                This could be because none of the bloq instances satisfied `pred` or none of
+                the bloqs have decompositions.
 
         """
         bb, _ = BloqBuilder.from_signature(self.signature)
@@ -348,13 +351,16 @@ def flatten_once(self, pred: Callable[[BloqInstance], bool]) -> 'CompositeBloq':
         bb._i = max(binst.i for binst in self.bloq_instances) + 1
 
         soq_map: List[Tuple[SoquetT, SoquetT]] = []
+        new_out_soqs: Tuple[SoquetT, ...] = ()
         did_work = False
         for binst, in_soqs, old_out_soqs in self.iter_bloqsoqs():
             in_soqs = _map_soqs(in_soqs, soq_map)  # update `in_soqs` from old to new.
-
             if pred(binst):
-                new_out_soqs = bb.add_from(binst.bloq.decompose_bloq(), **in_soqs)
-                did_work = True
+                try:
+                    new_out_soqs = bb.add_from(binst.bloq.decompose_bloq(), **in_soqs)
+                    did_work = True
+                except (DecomposeTypeError, DecomposeNotImplementedError):
+                    pass
             else:
                 # Since we took care to not re-use existing `binst.i` values for flattened
                 # bloqs, it is safe to call `bb._add_binst` with the old `binst` (and in
@@ -371,18 +377,8 @@ def flatten_once(self, pred: Callable[[BloqInstance], bool]) -> 'CompositeBloq':
         fsoqs = _map_soqs(self.final_soqs(), soq_map)
         return bb.finalize(**fsoqs)
 
-    def adjoint(self) -> 'CompositeBloq':
-        """Get a composite bloq which is the adjoint of this composite bloq.
-
-        The adjoint of a composite bloq is another composite bloq where the order of
-        operations is reversed and each subbloq is replaced with its adjoint.
-        """
-        from .adjoint import _adjoint_cbloq
-
-        return _adjoint_cbloq(self)
-
     def flatten(
-        self, pred: Callable[[BloqInstance], bool], max_depth: int = 1_000
+        self, pred: Callable[[BloqInstance], bool] = lambda binst: True, max_depth: int = 1_000
     ) -> 'CompositeBloq':
         """Recursively decompose and flatten subbloqs until none satisfy `pred`.
 
@@ -392,16 +388,13 @@ def flatten(
         Args:
             pred: A predicate that takes a bloq instance and returns True if it should
                 be decomposed and flattened or False if it should remain undecomposed.
-                All bloqs for which this callable returns True must support decomposition.
+                If the bloq does not have a decomposition, it will remain undecomposed.
+                By default, flatten as much as possible.
             max_depth: To avoid infinite recursion, give up after this many recursive steps.
 
         Returns:
             A new composite bloq where all recursive subbloqs matching `pred` have been
             decomposed and flattened.
-
-        Raises:
-            NotImplementedError: If `pred` returns True but the underlying bloq does not
-                support `decompose_bloq()`.
         """
         cbloq = self
         for _ in range(max_depth):
@@ -414,6 +407,16 @@ def flatten(
 
         return cbloq
 
+    def adjoint(self) -> 'CompositeBloq':
+        """Get a composite bloq which is the adjoint of this composite bloq.
+
+        The adjoint of a composite bloq is another composite bloq where the order of
+        operations is reversed and each subbloq is replaced with its adjoint.
+        """
+        from .adjoint import _adjoint_cbloq
+
+        return _adjoint_cbloq(self)
+
     @staticmethod
     def _debug_binst(g: nx.DiGraph, binst: BloqInstance) -> List[str]:
         """Helper method used in `debug_text`"""
@@ -454,6 +457,9 @@ def debug_text(self) -> str:
         delimited_gens = ('\n' + '-' * 20 + '\n').join(gen_texts)
         return delimited_gens
 
+    def __str__(self):
+        return f'CompositeBloq([{len(self.bloq_instances)} subbloqs...])'
+
 
 def _create_binst_graph(
     cxns: Iterable[Connection], nodes: Iterable[BloqInstance] = ()

qualtran/_infra/composite_bloq_test.py

@@ -30,7 +30,6 @@
     BloqInstance,
     CompositeBloq,
     Connection,
-    DecomposeTypeError,
     LeftDangle,
     Register,
     RightDangle,
@@ -504,13 +503,15 @@ def test_flatten():
     cbloq2 = cbloq.flatten_once(lambda binst: True)
     assert len(cbloq2.bloq_instances) == 5 * 2
 
-    with pytest.raises(DecomposeTypeError):
-        # Will keep trying to flatten non-decomposable things
-        cbloq.flatten(lambda x: True)
-
-    cbloq3 = cbloq.flatten(lambda binst: binst.bloq.supports_decompose_bloq())
+    cbloq3 = cbloq.flatten(lambda binst: True)
     assert len(cbloq3.bloq_instances) == 5 * 2
 
+    cbloq4 = cbloq.flatten(lambda binst: binst.bloq.supports_decompose_bloq())
+    assert len(cbloq4.bloq_instances) == 5 * 2
+
+    cbloq5 = cbloq.flatten()
+    assert len(cbloq5.bloq_instances) == 5 * 2
+
 
 def test_type_error():
     bb = BloqBuilder()