fix ctrl sys for CZ and CY: target register names of uncontrolled…

… and controlled versions do not match
quantumlib · Oct 29, 2024 · 6248dbe · 6248dbe
1 parent 5ea723a
commit 6248dbe
Show file tree

Hide file tree

Showing 3 changed files with 33 additions and 5 deletions.
diff --git a/qualtran/bloqs/basic_gates/y_gate.py b/qualtran/bloqs/basic_gates/y_gate.py
@@ -27,6 +27,7 @@
     ConnectionT,
     CtrlSpec,
     DecomposeTypeError,
+    QBit,
     Register,
     Signature,
 )
@@ -73,10 +74,15 @@ def my_tensors(
         ]
 
     def get_ctrl_system(self, ctrl_spec: 'CtrlSpec') -> Tuple['Bloq', 'AddControlledT']:
+        from qualtran.bloqs.bookkeeping import AutoPartition
         from qualtran.bloqs.mcmt.specialized_ctrl import get_ctrl_system_1bit_cv_from_bloqs
 
+        cy_wrapped = AutoPartition(
+            CYGate(), [(Register('ctrl', QBit()), ['ctrl']), (Register('q', QBit()), ['target'])]
+        )
+
         return get_ctrl_system_1bit_cv_from_bloqs(
-            self, ctrl_spec, current_ctrl_bit=None, bloq_with_ctrl=CYGate(), ctrl_reg_name='ctrl'
+            self, ctrl_spec, current_ctrl_bit=None, bloq_with_ctrl=cy_wrapped, ctrl_reg_name='ctrl'
         )
 
     def as_cirq_op(
@@ -164,10 +170,15 @@ def wire_symbol(
         raise ValueError(f"Unknown register {reg}.")
 
     def get_ctrl_system(self, ctrl_spec: 'CtrlSpec') -> Tuple['Bloq', 'AddControlledT']:
+        from qualtran.bloqs.bookkeeping import AutoPartition
         from qualtran.bloqs.mcmt.specialized_ctrl import get_ctrl_system_1bit_cv_from_bloqs
 
+        cy_wrapped = AutoPartition(
+            self, [(Register('ctrl', QBit()), ['ctrl']), (Register('q', QBit()), ['target'])]
+        )
+
         return get_ctrl_system_1bit_cv_from_bloqs(
-            self, ctrl_spec, current_ctrl_bit=1, bloq_with_ctrl=self, ctrl_reg_name='ctrl'
+            self, ctrl_spec, current_ctrl_bit=1, bloq_with_ctrl=cy_wrapped, ctrl_reg_name='ctrl'
         )
 
 

diff --git a/qualtran/bloqs/basic_gates/z_basis.py b/qualtran/bloqs/basic_gates/z_basis.py
@@ -254,10 +254,16 @@ def my_tensors(
         ]
 
     def get_ctrl_system(self, ctrl_spec: 'CtrlSpec') -> Tuple['Bloq', 'AddControlledT']:
+        from qualtran.bloqs.bookkeeping import AutoPartition
         from qualtran.bloqs.mcmt.specialized_ctrl import get_ctrl_system_1bit_cv_from_bloqs
 
+        cz = CZ()
+        cz_wrapped = AutoPartition(
+            cz, [(Register('ctrl', QBit()), ['q1']), (Register('q', QBit()), ['q2'])]
+        )
+
         return get_ctrl_system_1bit_cv_from_bloqs(
-            self, ctrl_spec, current_ctrl_bit=None, bloq_with_ctrl=CZ(), ctrl_reg_name='q1'
+            self, ctrl_spec, current_ctrl_bit=None, bloq_with_ctrl=cz_wrapped, ctrl_reg_name='ctrl'
         )
 
     def as_cirq_op(
@@ -332,10 +338,15 @@ def wire_symbol(self, reg: Optional[Register], idx: Tuple[int, ...] = tuple()) -
         raise ValueError(f'Unknown wire symbol register name: {reg.name}')
 
     def get_ctrl_system(self, ctrl_spec: 'CtrlSpec') -> Tuple['Bloq', 'AddControlledT']:
+        from qualtran.bloqs.bookkeeping import AutoPartition
         from qualtran.bloqs.mcmt.specialized_ctrl import get_ctrl_system_1bit_cv_from_bloqs
 
+        cz_wrapped = AutoPartition(
+            self, [(Register('ctrl', QBit()), ['q1']), (Register('q', QBit()), ['q2'])]
+        )
+
         return get_ctrl_system_1bit_cv_from_bloqs(
-            self, ctrl_spec, current_ctrl_bit=1, bloq_with_ctrl=self, ctrl_reg_name='q1'
+            self, ctrl_spec, current_ctrl_bit=1, bloq_with_ctrl=cz_wrapped, ctrl_reg_name='ctrl'
         )
 
 

diff --git a/qualtran/bloqs/mcmt/specialized_ctrl.py b/qualtran/bloqs/mcmt/specialized_ctrl.py
@@ -18,6 +18,7 @@
 import numpy as np
 
 from qualtran import Bloq, QBit, Register, Signature
+from qualtran.bloqs.bookkeeping import AutoPartition
 
 if TYPE_CHECKING:
     from qualtran import AddControlledT, BloqBuilder, CtrlSpec, SoquetT
@@ -180,7 +181,12 @@ def _adder(
 
             return [ctrl0], [ctrl1, *out_soqs]
 
-    return ctrl_bloq, _adder
+    def _unwrap(b):
+        if isinstance(b, AutoPartition):
+            return _unwrap(b.bloq)
+        return b
+
+    return _unwrap(ctrl_bloq), _adder
 
 
 def get_ctrl_system_1bit_cv(