imports and mypy

quantumlib · Oct 30, 2024 · e1dabca · e1dabca
1 parent 1147578
commit e1dabca
Show file tree

Hide file tree

Showing 3 changed files with 20 additions and 19 deletions.
diff --git a/qualtran/bloqs/mcmt/specialized_ctrl.py b/qualtran/bloqs/mcmt/specialized_ctrl.py
@@ -18,22 +18,11 @@
 import attrs
 import numpy as np
 
-from qualtran import (
-    AddControlledT,
-    Adjoint,
-    Bloq,
-    BloqBuilder,
-    CompositeBloq,
-    CtrlSpec,
-    QBit,
-    Register,
-    Signature,
-    SoquetT,
-)
+from qualtran import Adjoint, Bloq, BloqBuilder, CompositeBloq, QBit, Register, Signature
 from qualtran.bloqs.bookkeeping import AutoPartition
 
 if TYPE_CHECKING:
-    from qualtran import AddControlledT, BloqBuilder, CtrlSpec, SoquetT
+    from qualtran import AddControlledT, CtrlSpec, SoquetT
     from qualtran._infra.controlled import ControlBit
     from qualtran.resource_counting import BloqCountDictT, SympySymbolAllocator
 

diff --git a/qualtran/bloqs/mcmt/specialized_ctrl_test.py b/qualtran/bloqs/mcmt/specialized_ctrl_test.py
@@ -183,15 +183,20 @@ def test_bloq_with_controlled_bloq():
 def test_ctrl_adjoint():
     assert TestAtom('a').adjoint().controlled() == CTestAtom('a').adjoint()
 
-    _, sigma = TestAtom('g').adjoint().controlled(CtrlSpec(cvs=[1, 1])).call_graph(keep=_keep_and)
+    _, sigma = (
+        TestAtom('g')
+        .adjoint()
+        .controlled(ctrl_spec=CtrlSpec(cvs=[1, 1]))
+        .call_graph(keep=_keep_and)
+    )
     assert sigma == {And(): 1, And().adjoint(): 1, CTestAtom('g').adjoint(): 1}
 
     _, sigma = CTestAtom('c').adjoint().controlled().call_graph(keep=_keep_and)
     assert sigma == {And(): 1, And().adjoint(): 1, CTestAtom('c').adjoint(): 1}
 
     for cv in [0, 1]:
         assert (
-            AtomWithSpecializedControl().adjoint().controlled(CtrlSpec(cvs=(cv,)))
+            AtomWithSpecializedControl().adjoint().controlled(ctrl_spec=CtrlSpec(cvs=cv))
             == AtomWithSpecializedControl(cv=cv).adjoint()
         )
 

diff --git a/qualtran/bloqs/qubitization/qubitization_walk_operator.py b/qualtran/bloqs/qubitization/qubitization_walk_operator.py
@@ -33,8 +33,15 @@
 import cirq
 import numpy as np
 
-from qualtran import bloq_example, BloqBuilder, BloqDocSpec, Register, Signature, SoquetT
-from qualtran._infra.gate_with_registers import GateWithRegisters, total_bits
+from qualtran import (
+    bloq_example,
+    BloqBuilder,
+    BloqDocSpec,
+    GateWithRegisters,
+    Register,
+    Signature,
+    SoquetT,
+)
 from qualtran.bloqs.block_encoding.lcu_block_encoding import (
     BlackBoxPrepare,
     LCUBlockEncoding,
@@ -51,7 +58,7 @@
 
 
 @attrs.frozen(cache_hash=True)
-class QubitizationWalkOperator(GateWithRegisters):  # type: ignore[misc]
+class QubitizationWalkOperator(GateWithRegisters):
     r"""Construct a Szegedy Quantum Walk operator using LCU oracles SELECT and PREPARE.
 
     For a Hamiltonian $H = \sum_l w_l H_l$ (where coefficients $w_l > 0$ and $H_l$ are unitaries),
@@ -130,7 +137,7 @@ def build_composite_bloq(self, bb: 'BloqBuilder', **soqs: 'SoquetT') -> dict[str
         return soqs
 
     def _circuit_diagram_info_(self, args: cirq.CircuitDiagramInfoArgs) -> cirq.CircuitDiagramInfo:
-        wire_symbols = ['W'] * total_bits(self.signature)
+        wire_symbols = ['W'] * self.signature.n_qubits()
         return cirq.CircuitDiagramInfo(wire_symbols=wire_symbols)
 
     @cached_property