Add AutoPartition and BloqBuilder.add_and_partition to fit bloqs …

…together whose registers don't quite match (#1086)

* Add `BloqBuilder.add_and_partition`

* Separate `AutoPartition` bloq

* Add tests and documentation

* Fix mypy

* Update generalizers to pass through AutoPartition

* Update IPython notebook

* Address feedback

* Address feedback

* Fix generalizers_test

* Fix pylint

* Change flag name to `left_only` and update docs

---------

Co-authored-by: Matthew Harrigan <mpharrigan@google.com>

dev_tools/autogenerate-bloqs-notebooks-v2.py

@@ -192,6 +192,7 @@
             qualtran.bloqs.bookkeeping.split._SPLIT_DOC,
             qualtran.bloqs.bookkeeping.join._JOIN_DOC,
             qualtran.bloqs.bookkeeping.partition._PARTITION_DOC,
+            qualtran.bloqs.bookkeeping.auto_partition._AUTO_PARTITION_DOC,
             qualtran.bloqs.bookkeeping.cast._CAST_DOC,
         ],
     ),

qualtran/_infra/composite_bloq.py

@@ -966,6 +966,35 @@ def add_d(self, bloq: Bloq, **in_soqs: SoquetInT) -> Dict[str, SoquetT]:
         binst = BloqInstance(bloq, i=self._new_binst_i())
         return dict(self._add_binst(binst, in_soqs=in_soqs))
 
+    def add_and_partition(
+        self,
+        bloq: Bloq,
+        partitions: Sequence[Tuple[Register, Sequence[str]]],
+        left_only: bool = False,
+        **in_soqs: SoquetInT,
+    ):
+        """Add a new bloq instance to the compute graph by partitioning input and output soquets to
+        fit the signature of the bloq.
+
+        Args:
+            bloq: The bloq representing the operation to add.
+            partitions: A sequence of pairs specifying each register that the wrapped bloq should
+            accept and the register names from `bloq.signature.lefts()` that concatenate to form it.
+            left_only: If False, the output soquets will also follow `partition`.
+                Otherwise, the output soquets will follow `bloq.signature.rights()`.
+                This flag must be set to True if `bloq` does not have the same LEFT and RIGHT registers,
+                as is required for the bloq to be fully wrapped on the left and right.
+            **in_soqs: Keyword arguments mapping the new bloq's register names to input
+                `Soquet`s. This is likely the output soquets from a prior operation.
+
+        Returns:
+            A `Soquet` or an array thereof for each right (output) register ordered according to
+                `bloq.signature` or `partition`.
+        """
+        from qualtran.bloqs.bookkeeping.auto_partition import AutoPartition
+
+        return self.add(AutoPartition(bloq, partitions, left_only), **in_soqs)
+
     def add(self, bloq: Bloq, **in_soqs: SoquetInT):
         """Add a new bloq instance to the compute graph.
 

qualtran/_infra/composite_bloq_test.py

@@ -556,6 +556,24 @@ def test_t_complexity():
     assert TestParallelCombo().t_complexity().t == 3 * 100
 
 
+def test_add_and_partition():
+    from qualtran import Controlled, CtrlSpec
+    from qualtran.bloqs.basic_gates import Swap
+
+    bb = BloqBuilder()
+    bloq = Controlled(Swap(3), CtrlSpec(qdtypes=QUInt(4), cvs=0b0110))
+    a = bb.add_register_from_dtype('a', QAny(7))
+    b = bb.add_register_from_dtype('b', QAny(3))
+    assert a is not None
+    assert b is not None
+    a, b = bb.add_and_partition(
+        bloq, [(Register('a', QAny(7)), ['y', 'ctrl']), (Register('b', QAny(3)), ['x'])], a=a, b=b
+    )
+    cbloq = bb.finalize(a=a, b=b)
+    assert isinstance(cbloq, CompositeBloq)
+    assert len(cbloq.bloq_instances) == 1
+
+
 @pytest.mark.notebook
 def test_notebook():
     qlt_testing.execute_notebook('composite_bloq')

qualtran/bloqs/bookkeeping/__init__.py

@@ -15,6 +15,7 @@
 
 from qualtran.bloqs.bookkeeping.allocate import Allocate
 from qualtran.bloqs.bookkeeping.arbitrary_clifford import ArbitraryClifford
+from qualtran.bloqs.bookkeeping.auto_partition import AutoPartition
 from qualtran.bloqs.bookkeeping.cast import Cast
 from qualtran.bloqs.bookkeeping.free import Free
 from qualtran.bloqs.bookkeeping.join import Join

qualtran/bloqs/bookkeeping/auto_partition.py

@@ -0,0 +1,116 @@
+#  Copyright 2023 Google LLC
+#
+#  Licensed under the Apache License, Version 2.0 (the "License");
+#  you may not use this file except in compliance with the License.
+#  You may obtain a copy of the License at
+#
+#      https://www.apache.org/licenses/LICENSE-2.0
+#
+#  Unless required by applicable law or agreed to in writing, software
+#  distributed under the License is distributed on an "AS IS" BASIS,
+#  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+#  See the License for the specific language governing permissions and
+#  limitations under the License.
+from functools import cached_property
+from itertools import chain
+from typing import Dict, Sequence, Tuple
+
+from attrs import evolve, field, frozen
+
+from qualtran import (
+    Bloq,
+    bloq_example,
+    BloqBuilder,
+    BloqDocSpec,
+    QAny,
+    Register,
+    Side,
+    Signature,
+    SoquetT,
+)
+from qualtran.bloqs.bookkeeping.partition import Partition
+
+
+@frozen
+class AutoPartition(Bloq):
+    """Wrap a bloq with `Partition` to fit an alternative set of input and output registers
+       such that splits / joins can be avoided in diagrams.
+
+    Args:
+        bloq: The bloq to wrap.
+        partitions: A sequence of pairs specifying each register that the wrapped bloq should accept
+            and the names of registers from `bloq.signature.lefts()` that concatenate to form it.
+        left_only: If False, the output registers will also follow `partition`.
+            Otherwise, the output registers will follow `bloq.signature.rights()`.
+            This flag must be set to True if `bloq` does not have the same LEFT and RIGHT registers,
+            as is required for the bloq to be fully wrapped on the left and right.
+
+    Registers:
+        [user_spec]: The output registers of the wrapped bloq.
+    """
+
+    bloq: Bloq
+    partitions: Sequence[Tuple[Register, Sequence[str]]] = field(
+        converter=lambda s: tuple((r, tuple(rs)) for r, rs in s)
+    )
+    left_only: bool = False
+
+    def __attrs_post_init__(self):
+        regs = {r.name for r in self.bloq.signature.lefts()}
+        assert len(regs) == len(self.bloq.signature)
+        assert regs == {r for _, rs in self.partitions for r in rs}
+
+    @cached_property
+    def signature(self) -> Signature:
+        if self.left_only:
+            return Signature(
+                chain(
+                    (evolve(r, side=Side.LEFT) for r, _ in self.partitions),
+                    (evolve(r, side=Side.RIGHT) for r in self.bloq.signature.rights()),
+                )
+            )
+        else:
+            return Signature(r for r, _ in self.partitions)
+
+    def pretty_name(self) -> str:
+        return self.bloq.pretty_name()
+
+    def build_composite_bloq(self, bb: BloqBuilder, **soqs: SoquetT) -> Dict[str, SoquetT]:
+        parts: Dict[str, Partition] = dict()
+        in_regs: Dict[str, SoquetT] = dict()
+        for out_reg, bloq_regs in self.partitions:
+            part = Partition(
+                out_reg.bitsize, regs=tuple(self.bloq.signature.get_left(r) for r in bloq_regs)
+            )
+            parts[out_reg.name] = part
+            in_regs |= bb.add_d(part, x=soqs[out_reg.name])
+        bloq_out_regs = bb.add_d(self.bloq, **in_regs)
+        if self.left_only:
+            return bloq_out_regs
+
+        out_regs = {}
+        for soq_name in soqs.keys():
+            out_regs[soq_name] = bb.add(
+                parts[soq_name].adjoint(),
+                **{reg.name: bloq_out_regs[reg.name] for reg in parts[soq_name].signature.rights()},
+            )
+        return out_regs
+
+
+@bloq_example
+def _auto_partition() -> AutoPartition:
+    from qualtran import Controlled, CtrlSpec
+    from qualtran.bloqs.basic_gates import Swap
+
+    bloq = Controlled(Swap(1), CtrlSpec())
+    auto_partition = AutoPartition(
+        bloq, [(Register('x', QAny(2)), ['ctrl', 'x']), (Register('y', QAny(1)), ['y'])]
+    )
+    return auto_partition
+
+
+_AUTO_PARTITION_DOC = BloqDocSpec(
+    bloq_cls=AutoPartition,
+    import_line="from qualtran.bloqs.bookkeeping import AutoPartition",
+    examples=[_auto_partition],
+)

qualtran/bloqs/bookkeeping/auto_partition_test.py

@@ -0,0 +1,136 @@
+#  Copyright 2023 Google LLC
+#
+#  Licensed under the Apache License, Version 2.0 (the "License");
+#  you may not use this file except in compliance with the License.
+#  You may obtain a copy of the License at
+#
+#      https://www.apache.org/licenses/LICENSE-2.0
+#
+#  Unless required by applicable law or agreed to in writing, software
+#  distributed under the License is distributed on an "AS IS" BASIS,
+#  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+#  See the License for the specific language governing permissions and
+#  limitations under the License.
+import subprocess
+
+import pytest
+from attrs import evolve
+
+from qualtran.bloqs.bookkeeping.auto_partition import _auto_partition, AutoPartition
+
+
+def test_auto_partition(bloq_autotester):
+    bloq_autotester(_auto_partition)
+
+
+def test_auto_partition_input():
+    from qualtran import QAny, QBit, Register, Side
+
+    bloq = _auto_partition()
+
+    assert tuple(bloq.signature.lefts()) == (
+        Register('x', dtype=QAny(2)),
+        Register('y', dtype=QAny(1)),
+    )
+
+    assert tuple(bloq.signature.rights()) == (
+        Register('x', dtype=QAny(2)),
+        Register('y', dtype=QAny(1)),
+    )
+
+    x, y = bloq.call_classically(x=0b11, y=0b0)
+    assert x == 0b10
+    assert y == 0b1
+
+    x, y = bloq.call_classically(x=0b01, y=0b0)
+    assert x == 0b01
+    assert y == 0b0
+
+    bloq = evolve(bloq, left_only=True)
+
+    assert tuple(bloq.signature.lefts()) == (
+        Register('x', dtype=QAny(2), side=Side.LEFT),
+        Register('y', dtype=QAny(1), side=Side.LEFT),
+    )
+
+    assert tuple(bloq.signature.rights()) == (
+        Register('ctrl', dtype=QBit(), side=Side.RIGHT),
+        Register('x', dtype=QBit(), side=Side.RIGHT),
+        Register('y', dtype=QBit(), side=Side.RIGHT),
+    )
+
+    ctrl, x, y = bloq.call_classically(x=0b11, y=0b0)
+    assert ctrl == 0b1
+    assert x == 0b0
+    assert y == 0b1
+
+    ctrl, x, y = bloq.call_classically(x=0b01, y=0b0)
+    assert ctrl == 0b0
+    assert x == 0b1
+    assert y == 0b0
+
+
+def test_auto_partition_valid():
+    from qualtran import Controlled, CtrlSpec, QAny, QUInt, Register
+    from qualtran.bloqs.basic_gates import Swap
+
+    with pytest.raises(AssertionError):
+        bloq = Controlled(Swap(3), CtrlSpec(qdtypes=QUInt(4), cvs=0b0110))
+        bloq = AutoPartition(
+            bloq, [(Register('a', QAny(3)), ['y']), (Register('b', QAny(3)), ['x'])]
+        )
+
+
+def test_auto_partition_big():
+    from qualtran import Controlled, CtrlSpec, QAny, QUInt, Register, Side
+    from qualtran.bloqs.basic_gates import Swap
+
+    bloq = Controlled(Swap(3), CtrlSpec(qdtypes=QUInt(4), cvs=0b0110))
+    bloq = AutoPartition(
+        bloq, [(Register('a', QAny(7)), ['y', 'ctrl']), (Register('b', QAny(3)), ['x'])]
+    )
+
+    assert tuple(bloq.signature.lefts()) == (
+        Register('a', dtype=QAny(7)),
+        Register('b', dtype=QAny(3)),
+    )
+
+    assert tuple(bloq.signature.rights()) == (
+        Register('a', dtype=QAny(7)),
+        Register('b', dtype=QAny(3)),
+    )
+
+    a, b = bloq.call_classically(a=0b1010110, b=0b010)
+    assert a == 0b0100110
+    assert b == 0b101
+
+    a, b = bloq.call_classically(a=0b1010111, b=0b010)
+    assert a == 0b1010111
+    assert b == 0b010
+
+    bloq = evolve(bloq, left_only=True)
+
+    assert tuple(bloq.signature.lefts()) == (
+        Register('a', dtype=QAny(7), side=Side.LEFT),
+        Register('b', dtype=QAny(3), side=Side.LEFT),
+    )
+
+    assert tuple(bloq.signature.rights()) == (
+        Register('ctrl', dtype=QUInt(4), side=Side.RIGHT),
+        Register('x', dtype=QAny(3), side=Side.RIGHT),
+        Register('y', dtype=QAny(3), side=Side.RIGHT),
+    )
+
+    ctrl, x, y = bloq.call_classically(a=0b1010110, b=0b010)
+    assert ctrl == 0b0110
+    assert x == 0b101
+    assert y == 0b010
+
+    ctrl, x, y = bloq.call_classically(a=0b1010111, b=0b010)
+    assert ctrl == 0b0111
+    assert x == 0b010
+    assert y == 0b101
+
+
+def test_no_circular_import():
+    subprocess.check_call(['python', '-c', 'from qualtran.bloqs.bookkeeping import auto_partition'])