Changes required to please mypy with new SoquetT

Author: mpharrigan

qualtran/bloqs/basic_gates/rotation.py

@@ -220,8 +220,7 @@ def signature(self) -> 'Signature':
     def build_composite_bloq(self, bb: 'BloqBuilder', q: 'SoquetT') -> Dict[str, 'SoquetT']:
         from qualtran.bloqs.mcmt import And
 
-        q1, q2 = q  # type: ignore
-        (q1, q2), anc = bb.add(And(), ctrl=[q1, q2])
+        (q1, q2), anc = bb.add(And(), ctrl=q)
         anc = bb.add(ZPowGate(self.exponent, eps=self.eps), q=anc)
         (q1, q2) = bb.add(And().adjoint(), ctrl=[q1, q2], target=anc)
         return {'q': np.array([q1, q2])}

qualtran/bloqs/block_encoding/sparse_matrix.py

@@ -237,8 +237,7 @@ def build_composite_bloq(
         if is_symbolic(self.system_bitsize) or is_symbolic(self.row_oracle.num_nonzero):
             raise DecomposeTypeError(f"Cannot decompose symbolic {self=}")
 
-        assert not isinstance(ancilla, np.ndarray)
-        ancilla_bits = bb.split(ancilla)
+        ancilla_bits = bb.split(ancilla.item())
         q, l = ancilla_bits[0], bb.join(ancilla_bits[1:])
 
         l = bb.add(self.diffusion, target=l)

qualtran/bloqs/chemistry/pbc/first_quantization/select_and_prepare.py

@@ -153,11 +153,11 @@ def _reshape_reg(
         """
         # np.prod(()) returns a float (1.0), so take int
         size = int(np.prod(out_shape))
-        if isinstance(in_reg, np.ndarray):
+        if BloqBuilder.is_ndarray(in_reg):
             # split an array of bitsize qubits into flat list of qubits
             split_qubits = bb.split(bb.join(np.concatenate([bb.split(x) for x in in_reg.ravel()])))
         else:
-            split_qubits = bb.split(in_reg)
+            split_qubits = bb.split(in_reg.item())
         merged_qubits = np.array(
             [bb.join(split_qubits[i * bitsize : (i + 1) * bitsize]) for i in range(size)]
         )

qualtran/bloqs/chemistry/trotter/grid_ham/potential.py

@@ -27,7 +27,6 @@
     QAny,
     Register,
     Signature,
-    Soquet,
     SoquetT,
 )
 from qualtran._infra.data_types import BQUInt
@@ -93,7 +92,7 @@ def wire_symbol(
     def build_composite_bloq(
         self, bb: BloqBuilder, *, system_i: SoquetT, system_j: SoquetT
     ) -> Dict[str, SoquetT]:
-        if isinstance(system_i, Soquet) or isinstance(system_j, Soquet):
+        if not (BloqBuilder.is_ndarray(system_i) and BloqBuilder.is_ndarray(system_j)):
             raise ValueError("system_i and system_j must be numpy arrays of Soquet")
         # compute r_i - r_j
         # r_i + (-r_j), in practice we need to flip the sign bit, but this is just 3 cliffords.
@@ -227,7 +226,7 @@ def wire_symbol(
         return super().wire_symbol(reg, idx)
 
     def build_composite_bloq(self, bb: BloqBuilder, *, system: SoquetT) -> Dict[str, SoquetT]:
-        if isinstance(system, Soquet):
+        if not BloqBuilder.is_ndarray(system):
             raise ValueError("system must be a numpy array of Soquet")
         bitsize = (self.num_grid - 1).bit_length() + 1
         ij_pairs = np.triu_indices(self.num_elec, k=1)

qualtran/bloqs/data_loading/qroam_clean.py

@@ -520,8 +520,8 @@ def build_composite_bloq(self, bb: 'BloqBuilder', **soqs: 'SoquetT') -> Dict[str
         # Construct and return dictionary of final soquets.
         soqs |= {reg.name: soq for reg, soq in zip(self.control_registers, ctrl)}
         soqs |= {reg.name: soq for reg, soq in zip(self.selection_registers, selection)}
-        soqs |= {reg.name: soq.flat[1:] for reg, soq in zip(self.junk_registers, qrom_targets)}  # type: ignore[union-attr]
-        soqs |= {reg.name: soq.flat[0] for reg, soq in zip(self.target_registers, qrom_targets)}  # type: ignore[union-attr]
+        soqs |= {reg.name: soq.flat[1:] for reg, soq in zip(self.junk_registers, qrom_targets)}  # type: ignore[attr-defined]
+        soqs |= {reg.name: soq.flat[0] for reg, soq in zip(self.target_registers, qrom_targets)}  # type: ignore[attr-defined]
         return soqs
 
     def on_classical_vals(

qualtran/bloqs/for_testing/with_decomposition.py

@@ -13,16 +13,13 @@
 #  limitations under the License.
 
 from functools import cached_property
-from typing import Dict, TYPE_CHECKING
+from typing import Dict
 
 from attrs import frozen
 
-from qualtran import Bloq, BloqBuilder, Signature, Soquet
+from qualtran import Bloq, BloqBuilder, Signature, SoquetT
 from qualtran.bloqs.for_testing.atom import TestAtom
 
-if TYPE_CHECKING:
-    from qualtran import SoquetT
-
 
 @frozen
 class TestSerialCombo(Bloq):
@@ -47,7 +44,7 @@ def signature(self) -> Signature:
         return Signature.build(reg=3)
 
     def build_composite_bloq(self, bb: 'BloqBuilder', reg: 'SoquetT') -> Dict[str, 'SoquetT']:
-        assert isinstance(reg, Soquet)
+        assert BloqBuilder.is_single(reg)
         reg = bb.split(reg)
         for i in range(len(reg)):
             reg[i] = bb.add(TestAtom(), q=reg[i])

qualtran/bloqs/hamiltonian_simulation/hamiltonian_simulation_by_gqsp.py

@@ -19,7 +19,7 @@
 from attrs import field, frozen
 from numpy.typing import NDArray
 
-from qualtran import Bloq, bloq_example, BloqDocSpec, CtrlSpec, Signature, Soquet
+from qualtran import Bloq, bloq_example, BloqBuilder, BloqDocSpec, CtrlSpec, Signature, SoquetT
 from qualtran.bloqs.basic_gates.su2_rotation import SU2RotationGate
 from qualtran.bloqs.bookkeeping import Always
 from qualtran.bloqs.qsp.generalized_qsp import GeneralizedQSP
@@ -32,7 +32,6 @@
 from qualtran.symbolics import is_symbolic, Shaped, SymbolicFloat, SymbolicInt
 
 if TYPE_CHECKING:
-    from qualtran import BloqBuilder, SoquetT
     from qualtran.resource_counting import BloqCountDictT, SympySymbolAllocator
 
 
@@ -175,11 +174,11 @@ def build_composite_bloq(self, bb: 'BloqBuilder', **soqs: 'SoquetT') -> Dict[str
         soqs, state_prep_ancilla = self.__add_prepare(bb, soqs, state_prep_ancilla, adjoint=True)
 
         for soq in state_prep_ancilla.values():
-            if isinstance(soq, Soquet):
-                bb.free(soq)
+            if BloqBuilder.is_ndarray(soq):
+                for soq_element in soq.reshape(-1):
+                    bb.free(soq_element.item())
             else:
-                for soq_element in soq:
-                    bb.free(cast(Soquet, soq_element))
+                bb.free(soq.item())
 
         return soqs