fixes

Author: mpharrigan

qualtran/_infra/gate_with_registers.py

@@ -65,7 +65,7 @@ def split_qubits(
 
 def merge_qubits(
     registers: Iterable[Register],
-    **qubit_regs: Union['cirq.Qid', Sequence['cirq.Qid'], NDArray['cirq.Qid']],
+    **qubit_regs: Union['cirq.Qid', Sequence['cirq.Qid'], NDArray['cirq.Qid']],  # type: ignore[type-var]
 ) -> List['cirq.Qid']:
     """Merges the dictionary of appropriately shaped qubit arrays into a flat list of qubits."""
 
@@ -84,7 +84,7 @@ def merge_qubits(
     return ret
 
 
-def get_named_qubits(registers: Iterable[Register]) -> Dict[str, NDArray['cirq.Qid']]:
+def get_named_qubits(registers: Iterable[Register]) -> Dict[str, NDArray['cirq.Qid']]:  # type: ignore[type-var]
     """Returns a dictionary of appropriately shaped named qubit signature for input `signature`."""
 
     def _qubit_array(reg: Register):
@@ -327,7 +327,7 @@ def _num_qubits_(self) -> int:
         return total_bits(self.signature)
 
     def decompose_from_registers(
-        self, *, context: 'cirq.DecompositionContext', **quregs: NDArray['cirq.Qid']
+        self, *, context: 'cirq.DecompositionContext', **quregs: NDArray['cirq.Qid']  # type: ignore[type-var]
     ) -> 'cirq.OP_TREE':
         raise DecomposeNotImplementedError(f"{self} does not declare a decomposition.")
 
@@ -361,7 +361,7 @@ def on(self, *qubits) -> 'cirq.Operation':
         return cirq.Gate.on(self, *qubits)
 
     def on_registers(
-        self, **qubit_regs: Union['cirq.Qid', Sequence['cirq.Qid'], NDArray['cirq.Qid']]
+        self, **qubit_regs: Union['cirq.Qid', Sequence['cirq.Qid'], NDArray['cirq.Qid']]  # type: ignore[type-var]
     ) -> 'cirq.Operation':
         return self.on(*merge_qubits(self.signature, **qubit_regs))
 

qualtran/bloqs/arithmetic/addition.py

@@ -321,7 +321,7 @@ def with_registers(self, *new_registers: Union[int, Sequence[int]]):
         raise NotImplementedError("no need to implement with_registers.")
 
     def decompose_from_registers(
-        self, *, context: cirq.DecompositionContext, **quregs: NDArray[cirq.Qid]
+        self, *, context: cirq.DecompositionContext, **quregs
     ) -> cirq.OP_TREE:
         if not isinstance(self.bitsize, int):
             raise ValueError(f'Symbolic bitsize {self.bitsize} not supported')

qualtran/bloqs/arithmetic/comparison.py

@@ -253,7 +253,7 @@ def signature(self) -> Signature:
         )
 
     def decompose_from_registers(
-        self, *, context: cirq.DecompositionContext, **quregs: NDArray[cirq.Qid]
+        self, *, context: cirq.DecompositionContext, **quregs
     ) -> Iterator[cirq.OP_TREE]:
         x, y, ancilla = quregs['x'], quregs['y'], quregs['ancilla']
         x_msb, x_lsb = x
@@ -343,7 +343,7 @@ def signature(self) -> Signature:
         )
 
     def decompose_from_registers(
-        self, *, context: cirq.DecompositionContext, **quregs: NDArray[cirq.Qid]
+        self, *, context: cirq.DecompositionContext, **quregs
     ) -> Iterator[cirq.OP_TREE]:
         a = quregs['a']
         b = quregs['b']
@@ -508,7 +508,7 @@ def _decompose_via_tree(
         )
 
     def decompose_from_registers(
-        self, *, context: cirq.DecompositionContext, **quregs: NDArray[cirq.Qid]
+        self, *, context: cirq.DecompositionContext, **quregs
     ) -> Iterator[cirq.OP_TREE]:
         lhs, rhs, (target,) = list(quregs['x']), list(quregs['y']), quregs['target']
         input_qubits = set(lhs + rhs + [target])

qualtran/bloqs/data_loading/qrom.py

@@ -131,7 +131,7 @@ def _load_nth_data(
                 )
 
     def decompose_zero_selection(
-        self, context: cirq.DecompositionContext, **quregs: NDArray[cirq.Qid]
+        self, context: cirq.DecompositionContext, **quregs
     ) -> Iterator[cirq.OP_TREE]:
         controls = tuple(merge_qubits(self.control_registers, **quregs))
         target_regs = {reg.name: quregs[reg.name] for reg in self.target_registers}
@@ -154,7 +154,7 @@ def decompose_zero_selection(
             context.qubit_manager.qfree(list(junk.flatten()) + [and_target])
 
     def decompose_from_registers(
-        self, *, context: cirq.DecompositionContext, **quregs: NDArray[cirq.Qid]
+        self, *, context: cirq.DecompositionContext, **quregs
     ) -> cirq.OP_TREE:
         if self.has_data():
             return super().decompose_from_registers(context=context, **quregs)

qualtran/bloqs/mcmt/and_bloq.py

@@ -308,9 +308,9 @@ def __pow__(self, power: int) -> "Bloq":
 
     def _decompose_via_tree(
         self,
-        controls: NDArray[cirq.Qid],
+        controls: NDArray[cirq.Qid],  # type: ignore[type-var]
         control_values: Tuple[SymbolicInt, ...],
-        ancillas: NDArray[cirq.Qid],
+        ancillas: NDArray[cirq.Qid],  # type: ignore[type-var]
         target: cirq.Qid,
     ) -> Iterator[cirq.OP_TREE]:
         """Decomposes multi-controlled `And` in-terms of an `And` ladder of size #controls- 2."""
@@ -326,7 +326,7 @@ def _decompose_via_tree(
         yield from self._decompose_via_tree(new_controls, new_control_values, ancillas[1:], target)
 
     def decompose_from_registers(
-        self, *, context: cirq.DecompositionContext, **quregs: NDArray[cirq.Qid]
+        self, *, context: cirq.DecompositionContext, **quregs
     ) -> Iterator[cirq.OP_TREE]:
         control, ancilla, target = (
             quregs['ctrl'].flatten(),

qualtran/bloqs/multiplexers/selected_majorana_fermion.py

@@ -97,7 +97,7 @@ def extra_registers(self) -> Tuple[Register, ...]:
         return (Register('accumulator', QBit()),)
 
     def decompose_from_registers(
-        self, context: cirq.DecompositionContext, **quregs: NDArray[cirq.Qid]
+        self, context: cirq.DecompositionContext, **quregs
     ) -> Iterator[cirq.OP_TREE]:
         quregs['accumulator'] = np.array(context.qubit_manager.qalloc(1))
         control: Sequence['cirq.Qid'] = (

qualtran/bloqs/multiplexers/unary_iteration_bloq.py

@@ -507,7 +507,7 @@ def _break_early(
         return False
 
     def decompose_from_registers(
-        self, *, context: cirq.DecompositionContext, **quregs: NDArray[cirq.Qid]
+        self, *, context: cirq.DecompositionContext, **quregs
     ) -> cirq.OP_TREE:
         if total_bits(self.selection_registers) == 0 or self._break_early(
             (), 0, self.selection_registers[0].dtype.iteration_length_or_zero()

qualtran/bloqs/rotations/phase_gradient.py

@@ -200,7 +200,7 @@ def phase_dtype(self) -> QFxp:
         return QFxp(self.bitsize, self.bitsize)
 
     def decompose_from_registers(
-        self, *, context: cirq.DecompositionContext, **quregs: NDArray[cirq.Qid]
+        self, *, context: cirq.DecompositionContext, **quregs
     ) -> Iterator[cirq.OP_TREE]:
         if isinstance(self.bitsize, sympy.Expr):
             raise ValueError(f'Symbolic Bitsize not supported {self.bitsize}')
@@ -479,7 +479,7 @@ def scaled_val(self, x: int) -> int:
         return int(np.floor(result.astype(float) * 2**self.phase_dtype.bitsize) * sign)
 
     def decompose_from_registers(
-        self, *, context: cirq.DecompositionContext, **quregs: NDArray[cirq.Qid]
+        self, *, context: cirq.DecompositionContext, **quregs
     ) -> Iterator[cirq.OP_TREE]:
         if isinstance(self.gamma, sympy.Expr):
             raise ValueError(f'Symbolic gamma {self.gamma} not allowed')

qualtran/bloqs/rotations/programmable_rotation_gate_array.py

@@ -139,7 +139,7 @@ def signature(self) -> Signature:
         )
 
     def decompose_from_registers(
-        self, *, context: cirq.DecompositionContext, **quregs: NDArray[cirq.Qid]
+        self, *, context: cirq.DecompositionContext, **quregs
     ) -> Iterator[cirq.OP_TREE]:
         selection, kappa_load_target = quregs.pop('selection'), quregs.pop('kappa_load_target')
         rotations_target = quregs.pop('rotations_target')
@@ -208,7 +208,9 @@ def __init__(
         assert all(cirq.num_qubits(u) == self._target_bitsize for u in interleaved_unitaries)
         self._interleaved_unitaries = tuple(interleaved_unitaries)
 
-    def interleaved_unitary(self, index: int, **qubit_regs: NDArray[cirq.Qid]) -> cirq.Operation:
+    def interleaved_unitary(
+        self, index: int, **qubit_regs: NDArray[cirq.Qid]  # type: ignore[type-var]
+    ) -> cirq.Operation:
         return self._interleaved_unitaries[index].on(*qubit_regs['rotations_target'])
 
     @cached_property

qualtran/cirq_interop/_cirq_to_bloq.py

@@ -544,12 +544,12 @@ def cirq_optree_to_cbloq(
 
         reg_dtypes = [r.dtype for r in bloq.signature]
         # 3.1 Find input / output registers.
-        all_op_quregs: Dict[str, NDArray[_QReg]] = {
+        all_op_quregs: Dict[str, NDArray[_QReg]] = {  # type: ignore[type-var]
             k: np.apply_along_axis(_QReg, -1, *(v, reg_dtypes[i]))  # type: ignore
             for i, (k, v) in enumerate(split_qubits(bloq.signature, op.qubits).items())
         }
 
-        in_op_quregs: Dict[str, NDArray[_QReg]] = {
+        in_op_quregs: Dict[str, NDArray[_QReg]] = {  # type: ignore[type-var]
             reg.name: all_op_quregs[reg.name] for reg in bloq.signature.lefts()
         }
         # 3.2 Find input Soquets, by potentially allocating new Bloq registers corresponding to