Merge remote-tracking branch 'origin/main' into 2024-04/generic-costs-framework

Author: mpharrigan

.github/workflows/ci.yaml

@@ -99,3 +99,20 @@ jobs:
           pip install --no-deps -e .
       - run: |
           check/pylint
+  
+  mypy:
+    runs-on: ubuntu-latest
+    steps:
+      - uses: actions/checkout@v3
+        with:
+          fetch-depth: 0
+      - uses: actions/setup-python@v4
+        with:
+          python-version: "3.10"
+      - name: Install dependencies
+        run: |
+          python -m pip install --upgrade pip
+          pip install -r dev_tools/requirements/envs/dev.env.txt
+          pip install --no-deps -e .
+      - run: |
+          check/mypy

dev_tools/conf/mypy.ini

@@ -3,6 +3,7 @@ show_error_codes = true
 plugins = duet.typing, numpy.typing.mypy_plugin
 allow_redefinition = true
 check_untyped_defs = true
+
 # Disabling function override checking
 # Qualtran has many places where kwargs are used
 # with the intention to override in subclasses in ways mypy does not like

qualtran/_infra/bloq.py

@@ -27,9 +27,11 @@
 
     from qualtran import (
         AddControlledT,
+        Adjoint,
         BloqBuilder,
         CompositeBloq,
         CtrlSpec,
+        GateWithRegisters,
         Register,
         Signature,
         Soquet,

qualtran/_infra/composite_bloq.py

@@ -49,7 +49,6 @@
     import cirq
 
     from qualtran.cirq_interop._cirq_to_bloq import CirqQuregInT, CirqQuregT
-    from qualtran.cirq_interop.t_complexity_protocol import TComplexity
     from qualtran.resource_counting import BloqCountT, SympySymbolAllocator
     from qualtran.simulation.classical_sim import ClassicalValT
 
@@ -139,16 +138,19 @@ def as_cirq_op(
         """Return a cirq.CircuitOperation containing a cirq-exported version of this cbloq."""
         import cirq
 
-        circuit, out_quregs = self.to_cirq_circuit(qubit_manager=qubit_manager, **cirq_quregs)
+        circuit, out_quregs = self.to_cirq_circuit_and_quregs(
+            qubit_manager=qubit_manager, cirq_quregs=cirq_quregs
+        )
         return cirq.CircuitOperation(circuit), out_quregs
 
-    def to_cirq_circuit(
-        self, qubit_manager: Optional['cirq.QubitManager'] = None, **cirq_quregs: 'CirqQuregInT'
+    def to_cirq_circuit_and_quregs(
+        self, qubit_manager: Optional['cirq.QubitManager'] = None, **cirq_quregs
     ) -> Tuple['cirq.FrozenCircuit', Dict[str, 'CirqQuregT']]:
-        """Convert this CompositeBloq to a `cirq.Circuit`.
+        """Convert this CompositeBloq to a `cirq.Circuit` and output qubit registers.
 
         Args:
-            qubit_manager: A `cirq.QubitManager` to allocate new qubits.
+            qubit_manager: A `cirq.QubitManager` to allocate new qubits. If not provided,
+                uses `cirq.SimpleQubitManager()` by default.
             **cirq_quregs: Mapping from left register names to Cirq qubit arrays.
 
         Returns:
@@ -166,6 +168,30 @@ def to_cirq_circuit(
             self.signature, cirq_quregs, self._binst_graph, qubit_manager=qubit_manager
         )
 
+    def to_cirq_circuit(
+        self,
+        *,
+        qubit_manager: Optional['cirq.QubitManager'] = None,
+        cirq_quregs: Optional[Mapping[str, 'CirqQuregInT']] = None,
+    ) -> 'cirq.FrozenCircuit':
+        """Convert this CompositeBloq to a `cirq.Circuit`.
+
+        Args:
+            qubit_manager: A `cirq.QubitManager` to allocate new qubits. If not provided,
+                uses `cirq.SimpleQubitManager()` by default.
+            cirq_quregs: Mapping from left register names to Cirq qubit arrays. If not provided,
+                uses `get_named_qubits(self.signature.lefts())` by default.
+
+        Returns:
+            circuit: The cirq.FrozenCircuit version of this composite bloq.
+        """
+        from qualtran._infra.gate_with_registers import get_named_qubits
+
+        if cirq_quregs is None:
+            cirq_quregs = get_named_qubits(self.signature.lefts())
+
+        return self.to_cirq_circuit_and_quregs(qubit_manager=qubit_manager, **cirq_quregs)[0]
+
     @classmethod
     def from_cirq_circuit(cls, circuit: 'cirq.Circuit') -> 'CompositeBloq':
         """Construct a composite bloq from a Cirq circuit.

qualtran/_infra/composite_bloq_test.py

@@ -40,7 +40,6 @@
 )
 from qualtran._infra.composite_bloq import _create_binst_graph, _get_dangling_soquets
 from qualtran._infra.data_types import BoundedQUInt, QAny, QBit, QFxp, QUInt
-from qualtran._infra.gate_with_registers import get_named_qubits
 from qualtran.bloqs.basic_gates import CNOT, IntEffect, ZeroEffect
 from qualtran.bloqs.for_testing.atom import TestAtom, TestTwoBitOp
 from qualtran.bloqs.for_testing.many_registers import TestMultiTypedRegister, TestQFxp
@@ -164,7 +163,9 @@ def test_map_soqs():
 
 def test_bb_composite_bloq():
     cbloq_auto = TestTwoCNOT().decompose_bloq()
-    circuit, _ = cbloq_auto.to_cirq_circuit(q1=[cirq.LineQubit(1)], q2=[cirq.LineQubit(2)])
+    circuit, _ = cbloq_auto.to_cirq_circuit_and_quregs(
+        q1=[cirq.LineQubit(1)], q2=[cirq.LineQubit(2)]
+    )
     cirq.testing.assert_has_diagram(
         circuit,
         desired="""\
@@ -174,6 +175,16 @@ def test_bb_composite_bloq():
     """,
     )
 
+    circuit = cbloq_auto.to_cirq_circuit()
+    cirq.testing.assert_has_diagram(
+        circuit,
+        desired="""\
+q1: ───@───X───
+       │   │
+q2: ───X───@───
+    """,
+    )
+
 
 def test_bloq_builder():
     signature = Signature.build(x=1, y=1)
@@ -344,7 +355,7 @@ def test_complicated_target_register():
     # note: this includes the two `Dangling` generations.
     assert len(list(nx.topological_generations(binst_graph))) == 2 * 3 + 2
 
-    circuit, _ = cbloq.to_cirq_circuit(None, **get_named_qubits(bloq.signature.lefts()))
+    circuit = cbloq.to_cirq_circuit()
     cirq.testing.assert_has_diagram(
         circuit,
         """\

qualtran/_infra/controlled.py

@@ -445,7 +445,7 @@ def wire_symbol(self, reg: Register, idx: Tuple[int, ...] = tuple()) -> 'WireSym
         return self.ctrl_spec.wire_symbol(i, reg, idx)
 
     def adjoint(self) -> 'Bloq':
-        return self.subbloq.adjoint().controlled(self.ctrl_spec)
+        return self.subbloq.adjoint().controlled(ctrl_spec=self.ctrl_spec)
 
     def pretty_name(self) -> str:
         return f'C[{self.subbloq.pretty_name()}]'

qualtran/_infra/controlled_test.py

@@ -100,7 +100,7 @@ def test_ctrl_bloq_as_cirq_op():
 
     def _test_cirq_equivalence(bloq: Bloq, gate: cirq.Gate):
         left_quregs = get_named_qubits(bloq.signature.lefts())
-        circuit1, right_quregs = bloq.as_composite_bloq().to_cirq_circuit(None, **left_quregs)
+        circuit1 = bloq.as_composite_bloq().to_cirq_circuit(cirq_quregs=left_quregs)
         circuit2 = cirq.Circuit(
             gate.on(*merge_qubits(bloq.signature, **get_named_qubits(bloq.signature)))
         )
@@ -124,7 +124,7 @@ def _test_cirq_equivalence(bloq: Bloq, gate: cirq.Gate):
     bloq = Controlled(Swap(5), CtrlSpec(qdtypes=QUInt(4), cvs=0b0101))
     quregs = get_named_qubits(bloq.signature)
     ctrl, x, y = quregs['ctrl'], quregs['x'], quregs['y']
-    circuit1, _ = bloq.decompose_bloq().to_cirq_circuit(None, **quregs)
+    circuit1 = bloq.decompose_bloq().to_cirq_circuit(cirq_quregs=quregs)
     circuit2 = cirq.Circuit(
         cirq.SWAP(x[i], y[i]).controlled_by(*ctrl, control_values=[0, 1, 0, 1]) for i in range(5)
     )
@@ -415,7 +415,7 @@ def test_controlled_diagrams():
 q: ──────X^0.25───''',
     )
 
-    ctrl_0_gate = XPowGate(0.25).controlled(CtrlSpec(cvs=0))
+    ctrl_0_gate = XPowGate(0.25).controlled(ctrl_spec=CtrlSpec(cvs=0))
     cirq.testing.assert_has_diagram(
         cirq.Circuit(ctrl_0_gate.on_registers(**get_named_qubits(ctrl_0_gate.signature))),
         '''
@@ -424,7 +424,7 @@ def test_controlled_diagrams():
 q: ──────X^0.25───''',
     )
 
-    multi_ctrl_gate = XPowGate(0.25).controlled(CtrlSpec(cvs=[0, 1]))
+    multi_ctrl_gate = XPowGate(0.25).controlled(ctrl_spec=CtrlSpec(cvs=[0, 1]))
     cirq.testing.assert_has_diagram(
         cirq.Circuit(multi_ctrl_gate.on_registers(**get_named_qubits(multi_ctrl_gate.signature))),
         '''
@@ -435,7 +435,7 @@ def test_controlled_diagrams():
 q: ─────────X^0.25───''',
     )
 
-    ctrl_bloq = Swap(2).controlled(CtrlSpec(cvs=[0, 1]))
+    ctrl_bloq = Swap(2).controlled(ctrl_spec=CtrlSpec(cvs=[0, 1]))
     cirq.testing.assert_has_diagram(
         cirq.Circuit(ctrl_bloq.on_registers(**get_named_qubits(ctrl_bloq.signature))),
         '''

qualtran/_infra/gate_with_registers.py

@@ -15,6 +15,7 @@
 import abc
 from typing import (
     Any,
+    cast,
     Collection,
     Dict,
     Iterable,
@@ -27,6 +28,7 @@
     Union,
 )
 
+import attrs
 import cirq
 import numpy as np
 from numpy.typing import NDArray
@@ -38,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
 
@@ -357,8 +359,8 @@ def on_registers(
     ) -> cirq.Operation:
         return self.on(*merge_qubits(self.signature, **qubit_regs))
 
-    def __pow__(self, power: int) -> 'Bloq':
-        bloq = self if power > 0 else self.adjoint()
+    def __pow__(self, power: int) -> 'GateWithRegisters':
+        bloq = self if power > 0 else cast(GateWithRegisters, self.adjoint())
         if abs(power) == 1:
             return bloq
         if all(reg.side == Side.THRU for reg in self.signature):
@@ -381,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.
@@ -477,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 -
@@ -548,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/_infra/gate_with_registers_test.py

@@ -12,7 +12,7 @@
 #  See the License for the specific language governing permissions and
 #  limitations under the License.
 
-from typing import Dict, TYPE_CHECKING
+from typing import Dict, Iterator, TYPE_CHECKING
 
 import cirq
 import numpy as np
@@ -46,7 +46,7 @@ def signature(self) -> Signature:
         regs = Signature([r1, r2, r3])
         return regs
 
-    def decompose_from_registers(self, *, context, **quregs) -> cirq.OP_TREE:
+    def decompose_from_registers(self, *, context, **quregs) -> Iterator[cirq.OP_TREE]:
         yield cirq.H.on_each(quregs['r1'])
         yield cirq.X.on_each(quregs['r2'])
         yield cirq.X.on_each(quregs['r3'])

qualtran/bloqs/arithmetic/addition.py

@@ -14,7 +14,19 @@
 import itertools
 import math
 from functools import cached_property
-from typing import Any, Dict, Iterable, List, Optional, Sequence, Set, Tuple, TYPE_CHECKING, Union
+from typing import (
+    Any,
+    Dict,
+    Iterable,
+    Iterator,
+    List,
+    Optional,
+    Sequence,
+    Set,
+    Tuple,
+    TYPE_CHECKING,
+    Union,
+)
 
 import cirq
 import numpy as np
@@ -197,7 +209,7 @@ def _right_building_block(self, inp, out, anc, depth):
 
     def decompose_from_registers(
         self, *, context: cirq.DecompositionContext, **quregs: NDArray[cirq.Qid]  # type: ignore[type-var]
-    ) -> cirq.OP_TREE:
+    ) -> Iterator[cirq.OP_TREE]:
         # reverse the order of qubits for big endian-ness.
         input_bits = quregs['a'][::-1]
         output_bits = quregs['b'][::-1]