Fix Identity's ctrl_adder (#1781)

`Identity` is a special bloq: its controlled version is still an
identity, but with a larger register, ie. calling
`Identity(n).controlled(CtrlSpec(QAny(m)))` gets us `Identity(n+m)`

The old version of `ctrl_adder()` was using `Autopartition` which failed
during decomposition as `Autopartition` is not designed to merge
registers and pass them to the inner bloq.

This is fixed by this PR by replacing `Autopartition` with
`Partition(s)`.

This fix is functional but the resulting `show_bloq()` output is less
clean as the Partition steps are visible. An alternative (maybe better)
fix is to keep using `Autopartition` but mark the control registers as
`Unused()`, but It is weird conceptually to not use Identity(n+m) as the
controlled Identity.

This PR also adds a test covering this bug.

Author: petim0

qualtran/bloqs/basic_gates/identity.py

@@ -32,6 +32,7 @@
     Signature,
     SoquetT,
 )
+from qualtran.bloqs.bookkeeping import Partition
 from qualtran.drawing import Text, TextBox, WireSymbol
 from qualtran.symbolics import is_symbolic, SymbolicInt
 
@@ -116,11 +117,17 @@ def ctrl_adder(
             bb: 'BloqBuilder', ctrl_soqs: Sequence['SoquetT'], in_soqs: Dict[str, 'SoquetT']
         ) -> Tuple[Iterable['SoquetT'], Iterable['SoquetT']]:
             parts = [
-                (Register(f'ctrl_{i}', dtype=dtype, shape=shape), 'q')
+                Register(f"ctrl_{i}", dtype=dtype, shape=shape)
                 for i, (dtype, shape) in enumerate(ctrl_spec.activation_function_dtypes())
-            ] + [(reg, 'q') for reg in self.signature]
-            all_soqs = in_soqs | {f'ctrl_{i}': ctrl_soq for i, ctrl_soq in enumerate(ctrl_soqs)}
-            out_soqs = bb.add_and_partition(ctrl_I, partitions=parts, left_only=False, **all_soqs)
+            ] + [reg for reg in self.signature]
+
+            all_soqs = in_soqs | {f"ctrl_{i}": ctrl_soq for i, ctrl_soq in enumerate(ctrl_soqs)}
+
+            pratition = Partition(ctrl_I.signature.n_qubits(), regs=parts)
+            q = bb.add(pratition.adjoint(), **all_soqs)
+            q = bb.add(ctrl_I, q=q)
+            out_soqs = bb.add(pratition, x=q)
+
             return out_soqs[:-1], out_soqs[-1:]
 
         return ctrl_I, ctrl_adder

qualtran/bloqs/basic_gates/identity_test.py

@@ -16,7 +16,7 @@
 import pytest
 import sympy
 
-from qualtran import BloqBuilder
+from qualtran import BloqBuilder, CtrlSpec, QBit, Register
 from qualtran.bloqs.basic_gates import OneState
 from qualtran.bloqs.basic_gates.identity import _identity, _identity_n, _identity_symb, Identity
 from qualtran.simulation.classical_sim import (
@@ -106,6 +106,19 @@ def test_identity_controlled():
     assert Identity(n).controlled() == Identity(n + 1)
 
 
+def test_identity_ctrl_adder():
+
+    ctrl_I, ctrl_adder = Identity(1).get_ctrl_system(CtrlSpec())
+
+    bb = BloqBuilder()
+    ctrl0 = bb.add_register(Register("ctrl_0", QBit()))
+    q = bb.add_register(Register("q", QBit()))
+    assert ctrl0 is not None and q is not None
+    [ctrl_out], (out_reg,) = ctrl_adder(bb, ctrl_soqs=[ctrl0], in_soqs={"q": q})
+    composite = bb.finalize(ctrl_0=ctrl_out, q=out_reg)
+    composite.flatten()
+
+
 @pytest.mark.notebook
 def test_notebook():
     execute_notebook('identity')