Fix custom controlled versions of SelectOracles and QubitizationWalkOperators (#904)

* Fix custom controlled Select and Walk operators

* use mixin class instead

* `get_ctrl_spec` delegate to base class for other cases

* docstring typo

* docstring cleanup

* fix some type errors

* `type: ignore` the `controlled` mismatch errors

* remove superfluous `GateWithRegisters` inheritance

---------

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

Author: anurudhp

qualtran/_infra/gate_with_registers.py

@@ -28,6 +28,7 @@
     Union,
 )
 
+import attrs
 import cirq
 import numpy as np
 from numpy.typing import NDArray
@@ -39,7 +40,7 @@
 if TYPE_CHECKING:
     import quimb.tensor as qtn
 
-    from qualtran import CtrlSpec, SoquetT
+    from qualtran import AddControlledT, BloqBuilder, CtrlSpec, SoquetT
     from qualtran.cirq_interop import CirqQuregT
     from qualtran.drawing import WireSymbol
 
@@ -382,15 +383,6 @@ def _get_ctrl_spec(
         already accepts a `CtrlSpec` and simply returns it OR a Cirq-style API which accepts
         parameters expected by `cirq.Gate.controlled()` and converts them to a `CtrlSpec` object.
 
-        Users implementing custom `GateWithRegisters.controlled()` overrides can use this helper
-        to generate a CtrlSpec from the cirq-style API and thus easily support both Cirq & Bloq
-        APIs. For example
-
-        >>> class CustomGWR(GateWithRegisters):
-        >>>     def controlled(self, *args, **kwargs) -> 'Bloq':
-        >>>         ctrl_spec = self._get_ctrl_spec(*args, **kwargs)
-        >>>         # Use ctrl_spec to construct a controlled version of `self`.
-
         Args:
             num_controls: Cirq style API to specify control specification -
                 Total number of control qubits.
@@ -478,10 +470,6 @@ def controlled(
         bloq. Bloqs authors can declare their own, custom controlled versions by overriding
         `Bloq.get_ctrl_system` in the bloq.
 
-        If overriding the `GWR.controlled()` method directly, Bloq authors can use the
-        `self._get_ctrl_spec` helper to construct a `CtrlSpec` object from the input parameters of
-        `GWR.controlled()` and use it to return a custom controlled version of this Bloq.
-
 
         Args:
             num_controls: Cirq style API to specify control specification -
@@ -549,3 +537,62 @@ def _circuit_diagram_info_(self, args: cirq.CircuitDiagramInfoArgs) -> cirq.Circ
 
         wire_symbols[0] = self.__class__.__name__
         return cirq.CircuitDiagramInfo(wire_symbols=wire_symbols)
+
+
+class SpecializedSingleQubitControlledGate(GateWithRegisters):
+    """Add a specialized single-qubit controlled version of a Bloq.
+
+    `control_val` is an optional single-bit control. When `control_val` is provided,
+     the `control_registers` property should return a single named qubit register,
+     and otherwise return an empty tuple.
+
+    Example usage:
+
+        @attrs.frozen
+        class MyGate(SpecializedSingleQubitControlledGate):
+            control_val: Optional[int] = None
+
+            @property
+            def control_registers() -> Tuple[Register, ...]:
+                return () if self.control_val is None else (Register('control', QBit()),)
+    """
+
+    control_val: Optional[int]
+
+    @property
+    @abc.abstractmethod
+    def control_registers(self) -> Tuple[Register, ...]:
+        ...
+
+    def get_single_qubit_controlled_bloq(
+        self, control_val: int
+    ) -> 'SpecializedSingleQubitControlledGate':
+        """Override this to provide a custom controlled bloq"""
+        return attrs.evolve(self, control_val=control_val)  # type: ignore[misc]
+
+    def get_ctrl_system(
+        self, ctrl_spec: Optional['CtrlSpec'] = None
+    ) -> Tuple['Bloq', 'AddControlledT']:
+        if ctrl_spec is None:
+            ctrl_spec = CtrlSpec()
+
+        if self.control_val is None and ctrl_spec.shapes in [((),), ((1,),)]:
+            control_val = int(ctrl_spec.cvs[0].item())
+            cbloq = self.get_single_qubit_controlled_bloq(control_val)
+
+            if not hasattr(cbloq, 'control_registers'):
+                raise TypeError("{cbloq} should have attribute `control_registers`")
+
+            (ctrl_reg,) = cbloq.control_registers
+
+            def adder(
+                bb: 'BloqBuilder', ctrl_soqs: Sequence['SoquetT'], in_soqs: dict[str, 'SoquetT']
+            ) -> tuple[Iterable['SoquetT'], Iterable['SoquetT']]:
+                soqs = {ctrl_reg.name: ctrl_soqs[0]} | in_soqs
+                soqs = bb.add_d(cbloq, **soqs)
+                ctrl_soqs = [soqs.pop(ctrl_reg.name)]
+                return ctrl_soqs, soqs.values()
+
+            return cbloq, adder
+
+        return super().get_ctrl_system(ctrl_spec)

qualtran/bloqs/for_testing/random_select_and_prepare.py

@@ -12,15 +12,15 @@
 #  See the License for the specific language governing permissions and
 #  limitations under the License.
 from functools import cached_property
-from typing import Iterator, Optional, Sequence, Tuple
+from typing import Iterator, Optional, Tuple
 
-import attrs
 import cirq
 import numpy as np
 from attrs import frozen
 from numpy.typing import NDArray
 
 from qualtran import BloqBuilder, BoundedQUInt, QBit, Register, SoquetT
+from qualtran._infra.gate_with_registers import SpecializedSingleQubitControlledGate
 from qualtran.bloqs.for_testing.matrix_gate import MatrixGate
 from qualtran.bloqs.qubitization_walk_operator import QubitizationWalkOperator
 from qualtran.bloqs.select_and_prepare import PrepareOracle, SelectOracle
@@ -82,7 +82,7 @@ def alphas(self):
 
 
 @frozen
-class TestPauliSelectOracle(SelectOracle):
+class TestPauliSelectOracle(SpecializedSingleQubitControlledGate, SelectOracle):  # type: ignore[misc]
     r"""Paulis acting on $m$ qubits, controlled by an $n$-qubit register.
 
     Given $2^n$ multi-qubit-Paulis (acting on $m$ qubits) $U_j$,
@@ -132,33 +132,6 @@ def selection_registers(self) -> Tuple[Register, ...]:
     def target_registers(self) -> Tuple[Register, ...]:
         return (Register('target', BoundedQUInt(bitsize=self.target_bitsize)),)
 
-    def adjoint(self):
-        return self
-
-    def __pow__(self, power):
-        if abs(power) == 1:
-            return self
-        return NotImplemented
-
-    def controlled(
-        self,
-        num_controls: Optional[int] = None,
-        control_values=None,
-        control_qid_shape: Optional[Tuple[int, ...]] = None,
-    ) -> 'cirq.Gate':
-        if num_controls is None:
-            num_controls = 1
-        if control_values is None:
-            control_values = [1] * num_controls
-        if (
-            isinstance(control_values, Sequence)
-            and isinstance(control_values[0], int)
-            and len(control_values) == 1
-            and self.control_val is None
-        ):
-            return attrs.evolve(self, control_val=control_values[0])
-        raise NotImplementedError()
-
     def decompose_from_registers(
         self,
         *,
@@ -167,14 +140,12 @@ def decompose_from_registers(
         target: NDArray[cirq.Qid],  # type: ignore[type-var]
         **quregs: NDArray[cirq.Qid],  # type: ignore[type-var]
     ) -> Iterator[cirq.OP_TREE]:
-        if self.control_val is not None:
-            selection = np.concatenate([selection, quregs['control']])
-
         for cv, U in enumerate(self.select_unitaries):
             bits = tuple(map(int, bin(cv)[2:].zfill(self.select_bitsize)))[::-1]
+            op = U.on(*target).controlled_by(*selection, control_values=bits)
             if self.control_val is not None:
-                bits = (*bits, self.control_val)
-            yield U.on(*target).controlled_by(*selection, control_values=bits)
+                op = op.controlled_by(*quregs['control'], control_values=[self.control_val])
+            yield op
 
 
 def random_qubitization_walk_operator(

qualtran/bloqs/hubbard_model.py

@@ -47,15 +47,15 @@
 See the documentation for `PrepareHubbard` and `SelectHubbard` for details.
 """
 from functools import cached_property
-from typing import Collection, Iterator, Optional, Sequence, Tuple, TYPE_CHECKING, Union
+from typing import Iterator, Optional, Tuple, TYPE_CHECKING
 
 import attrs
 import cirq
 import numpy as np
 from numpy.typing import NDArray
 
 from qualtran import BoundedQUInt, QAny, QBit, Register, Signature
-from qualtran._infra.gate_with_registers import total_bits
+from qualtran._infra.gate_with_registers import SpecializedSingleQubitControlledGate, total_bits
 from qualtran.bloqs.arithmetic import AddConstantMod
 from qualtran.bloqs.basic_gates import CSwap
 from qualtran.bloqs.mcmt.and_bloq import MultiAnd
@@ -72,7 +72,7 @@
 
 
 @attrs.frozen
-class SelectHubbard(SelectOracle):
+class SelectHubbard(SpecializedSingleQubitControlledGate, SelectOracle):  # type: ignore[misc]
     r"""The SELECT operation optimized for the 2D Hubbard model.
 
     In contrast to the arbitrary chemistry Hamiltonian, we:
@@ -219,29 +219,6 @@ def decompose_from_registers(
             q_x=q_x, q_y=q_y, control=[*V, *control], target=target_qubits_for_apply_to_lth_gate
         )
 
-    def controlled(
-        self,
-        num_controls: Optional[int] = None,
-        control_values: Optional[
-            Union[cirq.ops.AbstractControlValues, Sequence[Union[int, Collection[int]]]]
-        ] = None,
-        control_qid_shape: Optional[Tuple[int, ...]] = None,
-    ) -> 'SelectHubbard':
-        if num_controls is None:
-            num_controls = 1
-        if control_values is None:
-            control_values = [1] * num_controls
-        if (
-            isinstance(control_values, Sequence)
-            and isinstance(control_values[0], int)
-            and len(control_values) == 1
-            and self.control_val is None
-        ):
-            return SelectHubbard(self.x_dim, self.y_dim, control_val=control_values[0])
-        raise NotImplementedError(
-            f'Cannot create a controlled version of {self} with control_values={control_values}.'
-        )
-
     def _circuit_diagram_info_(self, args: cirq.CircuitDiagramInfoArgs) -> cirq.CircuitDiagramInfo:
         info = super(SelectHubbard, self)._circuit_diagram_info_(args)
         if self.control_val is None:

qualtran/bloqs/hubbard_model_test.py

@@ -63,8 +63,6 @@ def test_hubbard_model_consistent_protocols():
         select_gate.controlled(num_controls=1, control_values=(0,)),
         select_op.controlled_by(cirq.q("control"), control_values=(0,)).gate,
     )
-    with pytest.raises(NotImplementedError, match="Cannot create a controlled version"):
-        _ = select_gate.controlled(num_controls=2)
 
     # Test diagrams
     expected_symbols = ['U', 'V', 'p_x', 'p_y', 'alpha', 'q_x', 'q_y', 'beta']

qualtran/bloqs/mean_estimation/mean_estimation_operator.py

@@ -13,14 +13,14 @@
 #  limitations under the License.
 
 from functools import cached_property
-from typing import Collection, Iterator, Optional, Sequence, Tuple, Union
+from typing import Iterator, Optional, Tuple
 
 import attrs
 import cirq
 from numpy.typing import NDArray
 
-from qualtran import GateWithRegisters, Register, Signature
-from qualtran._infra.gate_with_registers import total_bits
+from qualtran import CtrlSpec, Register, Signature
+from qualtran._infra.gate_with_registers import SpecializedSingleQubitControlledGate, total_bits
 from qualtran.bloqs.mean_estimation.complex_phase_oracle import ComplexPhaseOracle
 from qualtran.bloqs.reflection_using_prepare import ReflectionUsingPrepare
 from qualtran.bloqs.select_and_prepare import PrepareOracle, SelectOracle
@@ -63,7 +63,7 @@ def __attrs_post_init__(self):
 
 
 @attrs.frozen
-class MeanEstimationOperator(GateWithRegisters):
+class MeanEstimationOperator(SpecializedSingleQubitControlledGate):
     r"""Mean estimation operator $U=REFL_{p} ROT_{y}$ as per Sec 3.1 of arxiv.org:2208.07544.
 
     The MeanEstimationOperator (aka KO Operator) expects `CodeForRandomVariable` to specify the
@@ -82,21 +82,13 @@ class MeanEstimationOperator(GateWithRegisters):
     """
 
     code: CodeForRandomVariable
-    cv: Tuple[int, ...] = attrs.field(
-        converter=lambda v: (v,) if isinstance(v, int) else tuple(v), default=()
-    )
+    control_val: Optional[int] = None
     arctan_bitsize: int = 32
 
-    @cv.validator
-    def _validate_cv(self, attribute, value):
-        assert value in [(), (0,), (1,)]
-
     @cached_property
     def reflect(self) -> ReflectionUsingPrepare:
         return ReflectionUsingPrepare(
-            self.code.synthesizer,
-            control_val=None if self.cv == () else self.cv[0],
-            global_phase=-1,
+            self.code.synthesizer, global_phase=-1, control_val=self.control_val
         )
 
     @cached_property
@@ -126,36 +118,15 @@ def decompose_from_registers(
         yield self.select.on_registers(**select_reg)
         yield self.reflect.on_registers(**reflect_reg)
 
+    def get_single_qubit_controlled_bloq(self, control_val: int) -> 'MeanEstimationOperator':
+        c_encoder = self.code.encoder.controlled(ctrl_spec=CtrlSpec(cvs=control_val))
+        assert isinstance(c_encoder, SelectOracle)
+        c_code = attrs.evolve(self.code, encoder=c_encoder)
+        return attrs.evolve(self, code=c_code, control_val=control_val)
+
     def _circuit_diagram_info_(self, args: cirq.CircuitDiagramInfoArgs) -> cirq.CircuitDiagramInfo:
-        wire_symbols = [] if self.cv == () else [["(0)", "@"][self.cv[0]]]
+        wire_symbols = []
+        if self.control_val is not None:
+            wire_symbols.append("@" if self.control_val == 1 else "(0)")
         wire_symbols += ['U_ko'] * (total_bits(self.signature) - total_bits(self.control_registers))
         return cirq.CircuitDiagramInfo(wire_symbols=wire_symbols)
-
-    def controlled(
-        self,
-        num_controls: Optional[int] = None,
-        control_values: Optional[
-            Union[cirq.ops.AbstractControlValues, Sequence[Union[int, Collection[int]]]]
-        ] = None,
-        control_qid_shape: Optional[Tuple[int, ...]] = None,
-    ) -> 'MeanEstimationOperator':
-        if num_controls is None:
-            num_controls = 1
-        if control_values is None:
-            control_values = [1] * num_controls
-        if (
-            isinstance(control_values, Sequence)
-            and len(control_values) == 1
-            and isinstance(control_values[0], int)
-            and not self.cv
-        ):
-            c_select = self.code.encoder.controlled(control_values=control_values)
-            assert isinstance(c_select, SelectOracle)
-            return MeanEstimationOperator(
-                CodeForRandomVariable(encoder=c_select, synthesizer=self.code.synthesizer),
-                cv=self.cv + (control_values[0],),
-                arctan_bitsize=self.arctan_bitsize,
-            )
-        raise NotImplementedError(
-            f'Cannot create a controlled version of {self} with control_values={control_values}.'
-        )

qualtran/bloqs/mean_estimation/mean_estimation_operator_test.py

@@ -21,7 +21,11 @@
 from attrs import frozen
 
 from qualtran import BoundedQUInt, QAny, QBit, Register
-from qualtran._infra.gate_with_registers import get_named_qubits, total_bits
+from qualtran._infra.gate_with_registers import (
+    get_named_qubits,
+    SpecializedSingleQubitControlledGate,
+    total_bits,
+)
 from qualtran.bloqs.mean_estimation.mean_estimation_operator import (
     CodeForRandomVariable,
     MeanEstimationOperator,
@@ -51,7 +55,7 @@ def decompose_from_registers(  # type:ignore[override]
 
 
 @frozen
-class BernoulliEncoder(SelectOracle):
+class BernoulliEncoder(SpecializedSingleQubitControlledGate, SelectOracle):  # type: ignore[misc]
     r"""Encodes Bernoulli random variable y0/y1 as $Enc|ii..i>|0> = |ii..i>|y_{i}>$ where i=0/1."""
 
     p: float
@@ -86,10 +90,6 @@ def decompose_from_registers(  # type:ignore[override]
             if y1:
                 yield cirq.X(tq).controlled_by(*q, control_values=[1] * self.selection_bitsize)
 
-    def controlled(self, *args, **kwargs):
-        cv = kwargs['control_values'][0]
-        return BernoulliEncoder(self.p, self.y, self.selection_bitsize, self.target_bitsize, cv)
-
     @cached_property
     def mu(self) -> float:
         return self.p * self.y[1] + (1 - self.p) * self.y[0]
@@ -278,8 +278,6 @@ def test_mean_estimation_operator_consistent_protocols():
         mean_gate.controlled(num_controls=1, control_values=(0,)),
         op.controlled_by(cirq.q("control"), control_values=(0,)).gate,
     )
-    with pytest.raises(NotImplementedError, match="Cannot create a controlled version"):
-        _ = mean_gate.controlled(num_controls=2)
 
     # Test diagrams
     n_qubits = cirq.num_qubits(mean_gate)