Fix the final mypy --next errors (#5760)

Will push changes to github actions after this is in.

* `np.einsum` is missing a type signature, resulting in some type ignores (filed issue at numpy numpy/numpy#21978)
* In practice one can use numpy types like np.double for parameters, but currently we don't support that.  Filed #5758 to consider this.
* Numpy `tolist` does not return a type signature of `List`. wat?

Part of #3767

Author: dabacon

cirq-core/cirq/circuits/circuit_operation.py

@@ -20,13 +20,14 @@
 import math
 from typing import (
     Callable,
-    Mapping,
-    Sequence,
+    cast,
     Dict,
     FrozenSet,
     Iterator,
     List,
+    Mapping,
     Optional,
+    Sequence,
     Tuple,
     TYPE_CHECKING,
     Union,
@@ -78,7 +79,7 @@ class CircuitOperation(ops.Operation):
     def __init__(
         self,
         circuit: 'cirq.FrozenCircuit',
-        repetitions: int = 1,
+        repetitions: INT_TYPE = 1,
         qubit_map: Optional[Dict['cirq.Qid', 'cirq.Qid']] = None,
         measurement_key_map: Optional[Dict[str, str]] = None,
         param_resolver: Optional[study.ParamResolverOrSimilarType] = None,
@@ -790,4 +791,8 @@ def _resolve_parameters_(
         self, resolver: 'cirq.ParamResolver', recursive: bool
     ) -> 'cirq.CircuitOperation':
         resolved = self.with_params(resolver.param_dict, recursive)
-        return resolved.replace(repetitions=resolver.value_of(self.repetitions, recursive))
+        # repetitions can resolve to a float, but this is ok since constructor converts to
+        # nearby int.
+        return resolved.replace(
+            repetitions=resolver.value_of(cast('cirq.TParamVal', self.repetitions), recursive)
+        )

cirq-core/cirq/contrib/qasm_import/_parser.py

@@ -13,7 +13,7 @@
 # limitations under the License.
 import functools
 import operator
-from typing import Any, Callable, cast, Dict, Iterable, List, Optional, Sequence, Union
+from typing import Any, Callable, cast, Dict, Iterable, List, Optional, Union, TYPE_CHECKING
 
 import numpy as np
 import sympy
@@ -25,6 +25,10 @@
 from cirq.contrib.qasm_import.exception import QasmException
 
 
+if TYPE_CHECKING:
+    import cirq
+
+
 class Qasm:
     """Qasm stores the final result of the Qasm parsing."""
 
@@ -115,7 +119,10 @@ def on(
         # used as arguments, we generate reg_size GateOperations via iterating
         # through each qubit of the registers 0 to n-1 and use the same one
         # qubit from the "single-qubit registers" for each operation.
-        op_qubits = cast(Sequence[Sequence[ops.Qid]], functools.reduce(np.broadcast, args))
+        op_qubits = functools.reduce(
+            cast(Callable[[List['cirq.Qid'], List['cirq.Qid']], List['cirq.Qid']], np.broadcast),
+            args,
+        )
         for qubits in op_qubits:
             if isinstance(qubits, ops.Qid):
                 yield final_gate.on(qubits)

cirq-core/cirq/experiments/readout_confusion_matrix.py

@@ -175,9 +175,11 @@ def _get_vars(self, qubit_pattern: Sequence['cirq.Qid']) -> List[int]:
     def _confusion_matrix(self, qubits: Sequence['cirq.Qid']) -> np.ndarray:
         ein_input = []
         for qs, cm in zip(self.measure_qubits, self.confusion_matrices):
-            ein_input += [cm.reshape((2, 2) * len(qs)), self._get_vars(qs)]
+            ein_input.extend([cm.reshape((2, 2) * len(qs)), self._get_vars(qs)])
         ein_out = self._get_vars(qubits)
-        ret = np.einsum(*ein_input, ein_out).reshape((2 ** len(qubits),) * 2)
+
+        # TODO(#5757): remote type ignore when numpy has proper override signature.
+        ret = np.einsum(*ein_input, ein_out).reshape((2 ** len(qubits),) * 2)  # type: ignore
         return ret / ret.sum(axis=1)
 
     def confusion_matrix(self, qubits: Optional[Sequence['cirq.Qid']] = None) -> np.ndarray:

cirq-core/cirq/linalg/decompositions.py

@@ -23,7 +23,6 @@
     Iterable,
     List,
     Optional,
-    Sequence,
     Set,
     Tuple,
     TYPE_CHECKING,
@@ -637,7 +636,7 @@ def scatter_plot_normalized_kak_interaction_coefficients(
         ax = fig.add_subplot(1, 1, 1, projection='3d')
 
     def coord_transform(
-        pts: Sequence[Tuple[float, float, float]]
+        pts: Union[List[Tuple[int, int, int]], np.ndarray]
     ) -> Tuple[Iterable[float], Iterable[float], Iterable[float]]:
         if len(pts) == 0:
             return [], [], []

cirq-core/cirq/linalg/predicates.py

@@ -161,7 +161,10 @@ def is_cptp(*, kraus_ops: Sequence[np.ndarray], rtol: float = 1e-5, atol: float
         atol: The absolute tolerance on equality.
     """
     sum_ndarray = cast(np.ndarray, sum(matrix.T.conj() @ matrix for matrix in kraus_ops))
-    return np.allclose(sum_ndarray, np.eye(*sum_ndarray.shape), rtol=rtol, atol=atol)
+    # Explicitly pull out shapes and don't use tuple to avoid confusing numpy type overrides.
+    return np.allclose(
+        sum_ndarray, np.eye(sum_ndarray.shape[0], sum_ndarray.shape[1]), rtol=rtol, atol=atol
+    )
 
 
 def matrix_commutes(

cirq-core/cirq/linalg/transformations.py

@@ -151,6 +151,7 @@ def targeted_left_multiply(
 
     all_indices = set(input_indices + data_indices + tuple(output_indices))
 
+    # TODO(#5757): remote type ignore when numpy has proper override signature.
     return np.einsum(
         left_matrix,
         input_indices,
@@ -164,7 +165,7 @@ def targeted_left_multiply(
         # And this is workaround for *another* bug!
         # Supposed to be able to just say 'old=old'.
         **({'out': out} if out is not None else {}),
-    )
+    )  # type: ignore
 
 
 def targeted_conjugate_about(
@@ -333,7 +334,8 @@ def partial_trace(tensor: np.ndarray, keep_indices: Sequence[int]) -> np.ndarray
     keep_map = dict(zip(keep_indices, sorted(keep_indices)))
     left_indices = [keep_map[i] if i in keep_set else i for i in range(ndim)]
     right_indices = [ndim + i if i in keep_set else i for i in left_indices]
-    return np.einsum(tensor, left_indices + right_indices)
+    # TODO(#5757): remote type ignore when numpy has proper override signature.
+    return np.einsum(tensor, left_indices + right_indices)  # type: ignore
 
 
 class EntangledStateError(ValueError):

cirq-core/cirq/qis/states.py

@@ -677,13 +677,14 @@ def density_matrix_from_state_vector(
     sum_inds = np.array(range(n_qubits))
     sum_inds[indices] += n_qubits
 
+    # TODO(#5757): remote type ignore when numpy has proper override signature.
     rho = np.einsum(
         state_vector,
         list(range(n_qubits)),
         np.conj(state_vector),
-        sum_inds.tolist(),
-        indices + sum_inds[indices].tolist(),
-    )
+        cast(List, sum_inds.tolist()),
+        indices + cast(List, sum_inds[indices].tolist()),
+    )  # type: ignore
     new_shape = np.prod([shape[i] for i in indices], dtype=np.int64)
 
     return rho.reshape((new_shape, new_shape))

cirq-core/cirq/sim/state_vector.py

@@ -352,5 +352,4 @@ def _probs(state: np.ndarray, indices: Sequence[int], qid_shape: Tuple[int, ...]
         probs = np.sum(probs, axis=tuple(range(1, len(probs.shape))))
 
     # To deal with rounding issues, ensure that the probabilities sum to 1.
-    probs /= np.sum(probs)
-    return probs
+    return probs / np.sum(probs)

cirq-core/cirq/study/resolver.py

@@ -14,7 +14,7 @@
 
 """Resolves ParameterValues to assigned values."""
 import numbers
-from typing import Any, Dict, Iterator, Mapping, Optional, TYPE_CHECKING, Union, cast
+from typing import Any, cast, Dict, Iterator, Mapping, Optional, TYPE_CHECKING, Union
 
 import numpy as np
 import sympy

cirq-google/cirq_google/serialization/circuit_serializer_test.py

@@ -53,7 +53,7 @@ def circuit_proto(json: Dict, qubits: List[str]):
 
 X_PROTO = op_proto({'xpowgate': {'exponent': {'float_value': 1.0}}, 'qubit_constant_index': [0]})
 
-
+# TODO(#5758): Add support for numpy types to `TParamVal`.
 OPERATIONS = [
     (cirq.X(Q0), X_PROTO),
     (
@@ -69,11 +69,11 @@ def circuit_proto(json: Dict, qubits: List[str]):
         op_proto({'xpowgate': {'exponent': {'float_value': 0.125}}, 'qubit_constant_index': [0]}),
     ),
     (
-        cirq.XPowGate(exponent=np.double(0.125))(Q1),
+        cirq.XPowGate(exponent=np.double(0.125))(Q1),  # type: ignore
         op_proto({'xpowgate': {'exponent': {'float_value': 0.125}}, 'qubit_constant_index': [0]}),
     ),
     (
-        cirq.XPowGate(exponent=np.short(1))(Q1),
+        cirq.XPowGate(exponent=np.short(1))(Q1),  # type: ignore
         op_proto({'xpowgate': {'exponent': {'float_value': 1.0}}, 'qubit_constant_index': [0]}),
     ),
     (

examples/direct_fidelity_estimation.py

@@ -419,7 +419,7 @@ def direct_fidelity_estimation(
         measured_pauli_traces = pauli_traces
     else:
         # Otherwise, randomly sample as per probability.
-        measured_pauli_traces = np.random.choice(pauli_traces, size=len(pauli_traces), p=p)
+        measured_pauli_traces = np.random.choice(pauli_traces, size=len(pauli_traces), p=p).tolist()
 
     trial_results: List[Result] = []
     for pauli_trace in measured_pauli_traces: