cirq-core 1.5.0.dev20250409225226__py3-none-any.whl → 1.6.0__py3-none-any.whl

cirq/sim/simulation_product_state_test.py

@@ -11,7 +11,10 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
-from typing import Any, Dict, Optional, Sequence
+
+from __future__ import annotations
+
+from typing import Any, Sequence
 
 import cirq
 
@@ -42,7 +45,7 @@ class EmptySimulationState(cirq.SimulationState):
         super().__init__(state=EmptyQuantumState(), qubits=qubits, classical_data=classical_data)
 
     def _act_on_fallback_(
-        self, action: Any, qubits: Sequence['cirq.Qid'], allow_decompose: bool = True
+        self, action: Any, qubits: Sequence[cirq.Qid], allow_decompose: bool = True
     ) -> bool:
         return True
 
@@ -52,9 +55,9 @@ qs2 = cirq.LineQubit.range(2)
 
 
 def create_container(
-    qubits: Sequence['cirq.Qid'], split_untangled_states=True
+    qubits: Sequence[cirq.Qid], split_untangled_states=True
 ) -> cirq.SimulationProductState[EmptySimulationState]:
-    state_map: Dict[Optional['cirq.Qid'], EmptySimulationState] = {}
+    state_map: dict[cirq.Qid | None, EmptySimulationState] = {}
     log = cirq.ClassicalDataDictionaryStore()
     if split_untangled_states:
         for q in reversed(qubits):
@@ -70,7 +73,7 @@ def create_container(
     )
 
 
-def test_entanglement_causes_join():
+def test_entanglement_causes_join() -> None:
     state = create_container(qs2)
     assert len(set(state.values())) == 3
     state.apply_operation(cirq.CNOT(q0, q1))
@@ -79,7 +82,7 @@ def test_entanglement_causes_join():
     assert state[None] is not state[q0]
 
 
-def test_subcircuit_entanglement_causes_join():
+def test_subcircuit_entanglement_causes_join() -> None:
     state = create_container(qs2)
     assert len(set(state.values())) == 3
     state.apply_operation(cirq.CircuitOperation(cirq.FrozenCircuit(cirq.CNOT(q0, q1))))
@@ -87,7 +90,7 @@ def test_subcircuit_entanglement_causes_join():
     assert state[q0] is state[q1]
 
 
-def test_subcircuit_entanglement_causes_join_in_subset():
+def test_subcircuit_entanglement_causes_join_in_subset() -> None:
     state = create_container(qs3)
     assert len(set(state.values())) == 4
     state.apply_operation(cirq.CircuitOperation(cirq.FrozenCircuit(cirq.CNOT(q0, q1))))
@@ -98,7 +101,7 @@ def test_subcircuit_entanglement_causes_join_in_subset():
     assert state[q0] is state[q1] is state[q2]
 
 
-def test_identity_does_not_join():
+def test_identity_does_not_join() -> None:
     state = create_container(qs2)
     assert len(set(state.values())) == 3
     state.apply_operation(cirq.IdentityGate(2)(q0, q1))
@@ -107,7 +110,7 @@ def test_identity_does_not_join():
     assert state[q0] is not state[None]
 
 
-def test_identity_fallback_does_not_join():
+def test_identity_fallback_does_not_join() -> None:
     state = create_container(qs2)
     assert len(set(state.values())) == 3
     state._act_on_fallback_(cirq.I, (q0, q1))
@@ -116,7 +119,7 @@ def test_identity_fallback_does_not_join():
     assert state[q0] is not state[None]
 
 
-def test_subcircuit_identity_does_not_join():
+def test_subcircuit_identity_does_not_join() -> None:
     state = create_container(qs2)
     assert len(set(state.values())) == 3
     state.apply_operation(cirq.CircuitOperation(cirq.FrozenCircuit(cirq.IdentityGate(2)(q0, q1))))
@@ -124,7 +127,7 @@ def test_subcircuit_identity_does_not_join():
     assert state[q0] is not state[q1]
 
 
-def test_measurement_causes_split():
+def test_measurement_causes_split() -> None:
     state = create_container(qs2)
     state.apply_operation(cirq.CNOT(q0, q1))
     assert len(set(state.values())) == 2
@@ -134,7 +137,7 @@ def test_measurement_causes_split():
     assert state[q0] is not state[None]
 
 
-def test_subcircuit_measurement_causes_split():
+def test_subcircuit_measurement_causes_split() -> None:
     state = create_container(qs2)
     state.apply_operation(cirq.CNOT(q0, q1))
     assert len(set(state.values())) == 2
@@ -143,7 +146,7 @@ def test_subcircuit_measurement_causes_split():
     assert state[q0] is not state[q1]
 
 
-def test_subcircuit_measurement_causes_split_in_subset():
+def test_subcircuit_measurement_causes_split_in_subset() -> None:
     state = create_container(qs3)
     state.apply_operation(cirq.CNOT(q0, q1))
     state.apply_operation(cirq.CNOT(q0, q2))
@@ -158,7 +161,7 @@ def test_subcircuit_measurement_causes_split_in_subset():
     assert state[q1] is not state[q2]
 
 
-def test_reset_causes_split():
+def test_reset_causes_split() -> None:
     state = create_container(qs2)
     state.apply_operation(cirq.CNOT(q0, q1))
     assert len(set(state.values())) == 2
@@ -168,7 +171,7 @@ def test_reset_causes_split():
     assert state[q0] is not state[None]
 
 
-def test_measurement_does_not_split_if_disabled():
+def test_measurement_does_not_split_if_disabled() -> None:
     state = create_container(qs2, False)
     state.apply_operation(cirq.CNOT(q0, q1))
     assert len(set(state.values())) == 1
@@ -178,7 +181,7 @@ def test_measurement_does_not_split_if_disabled():
     assert state[None] is state[q0]
 
 
-def test_reset_does_not_split_if_disabled():
+def test_reset_does_not_split_if_disabled() -> None:
     state = create_container(qs2, False)
     state.apply_operation(cirq.CNOT(q0, q1))
     assert len(set(state.values())) == 1
@@ -188,7 +191,7 @@ def test_reset_does_not_split_if_disabled():
     assert state[None] is state[q0]
 
 
-def test_measurement_of_all_qubits_causes_split():
+def test_measurement_of_all_qubits_causes_split() -> None:
     state = create_container(qs2)
     state.apply_operation(cirq.CNOT(q0, q1))
     assert len(set(state.values())) == 2
@@ -198,7 +201,7 @@ def test_measurement_of_all_qubits_causes_split():
     assert state[q0] is not state[None]
 
 
-def test_measurement_in_single_qubit_circuit_passes():
+def test_measurement_in_single_qubit_circuit_passes() -> None:
     state = create_container([q0])
     assert len(set(state.values())) == 2
     state.apply_operation(cirq.measure(q0))
@@ -206,25 +209,25 @@ def test_measurement_in_single_qubit_circuit_passes():
     assert state[q0] is not state[None]
 
 
-def test_reorder_succeeds():
+def test_reorder_succeeds() -> None:
     state = create_container(qs2, False)
     reordered = state[q0].transpose_to_qubit_order([q1, q0])
     assert reordered.qubits == (q1, q0)
 
 
-def test_copy_succeeds():
+def test_copy_succeeds() -> None:
     state = create_container(qs2, False)
     copied = state[q0].copy()
     assert copied.qubits == (q0, q1)
 
 
-def test_merge_succeeds():
+def test_merge_succeeds() -> None:
     state = create_container(qs2, False)
     merged = state.create_merged_state()
     assert merged.qubits == (q0, q1)
 
 
-def test_swap_does_not_merge():
+def test_swap_does_not_merge() -> None:
     state = create_container(qs2)
     old_q0 = state[q0]
     old_q1 = state[q1]
@@ -238,14 +241,14 @@ def test_swap_does_not_merge():
     assert state[q1].qubits == (q1,)
 
 
-def test_half_swap_does_merge():
+def test_half_swap_does_merge() -> None:
     state = create_container(qs2)
     state.apply_operation(cirq.SWAP(q0, q1) ** 0.5)
     assert len(set(state.values())) == 2
     assert state[q0] is state[q1]
 
 
-def test_swap_after_entangle_reorders():
+def test_swap_after_entangle_reorders() -> None:
     state = create_container(qs2)
     state.apply_operation(cirq.CX(q0, q1))
     assert len(set(state.values())) == 2
@@ -256,7 +259,7 @@ def test_swap_after_entangle_reorders():
     assert state[q0].qubits == (q1, q0)
 
 
-def test_act_on_gate_does_not_join():
+def test_act_on_gate_does_not_join() -> None:
     state = create_container(qs2)
     assert len(set(state.values())) == 3
     cirq.act_on(cirq.X, state, [q0])
@@ -265,7 +268,7 @@ def test_act_on_gate_does_not_join():
     assert state[q0] is not state[None]
 
 
-def test_field_getters():
+def test_field_getters() -> None:
     state = create_container(qs2)
     assert state.sim_states.keys() == set(qs2) | {None}
     assert state.split_untangled_states

cirq/sim/simulation_state.py

@@ -11,23 +11,14 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
+
 """Objects and methods for acting efficiently on a state tensor."""
+
+from __future__ import annotations
+
 import abc
 import copy
-from typing import (
-    Any,
-    cast,
-    Dict,
-    Generic,
-    Iterator,
-    List,
-    Optional,
-    Sequence,
-    Set,
-    Tuple,
-    TYPE_CHECKING,
-    TypeVar,
-)
+from typing import Any, cast, Generic, Iterator, Sequence, TYPE_CHECKING, TypeVar
 
 import numpy as np
 from typing_extensions import Self
@@ -48,9 +39,9 @@ class SimulationState(SimulationStateBase, Generic[TState], metaclass=abc.ABCMet
         self,
         *,
         state: TState,
-        prng: Optional[np.random.RandomState] = None,
-        qubits: Optional[Sequence['cirq.Qid']] = None,
-        classical_data: Optional['cirq.ClassicalDataStore'] = None,
+        prng: np.random.RandomState | None = None,
+        qubits: Sequence[cirq.Qid] | None = None,
+        classical_data: cirq.ClassicalDataStore | None = None,
     ):
         """Inits SimulationState.
 
@@ -79,10 +70,10 @@ class SimulationState(SimulationStateBase, Generic[TState], metaclass=abc.ABCMet
 
     def measure(
         self,
-        qubits: Sequence['cirq.Qid'],
+        qubits: Sequence[cirq.Qid],
         key: str,
         invert_mask: Sequence[bool],
-        confusion_map: Dict[Tuple[int, ...], np.ndarray],
+        confusion_map: dict[tuple[int, ...], np.ndarray],
     ):
         """Measures the qubits and records to `log_of_measurement_results`.
 
@@ -106,10 +97,10 @@ class SimulationState(SimulationStateBase, Generic[TState], metaclass=abc.ABCMet
             value.MeasurementKey.parse_serialized(key), corrected, qubits
         )
 
-    def get_axes(self, qubits: Sequence['cirq.Qid']) -> List[int]:
+    def get_axes(self, qubits: Sequence[cirq.Qid]) -> list[int]:
         return [self.qubit_map[q] for q in qubits]
 
-    def _perform_measurement(self, qubits: Sequence['cirq.Qid']) -> List[int]:
+    def _perform_measurement(self, qubits: Sequence[cirq.Qid]) -> list[int]:
         """Delegates the call to measure the `QuantumStateRepresentation`."""
         if self._state is not None:
             return self._state.measure(self.get_axes(qubits), self.prng)
@@ -117,9 +108,9 @@ class SimulationState(SimulationStateBase, Generic[TState], metaclass=abc.ABCMet
 
     def _confuse_result(
         self,
-        bits: List[int],
-        qubits: Sequence['cirq.Qid'],
-        confusion_map: Dict[Tuple[int, ...], np.ndarray],
+        bits: list[int],
+        qubits: Sequence[cirq.Qid],
+        confusion_map: dict[tuple[int, ...], np.ndarray],
     ):
         """Applies confusion matrices to measured results.
 
@@ -138,9 +129,9 @@ class SimulationState(SimulationStateBase, Generic[TState], metaclass=abc.ABCMet
 
     def sample(
         self,
-        qubits: Sequence['cirq.Qid'],
+        qubits: Sequence[cirq.Qid],
         repetitions: int = 1,
-        seed: 'cirq.RANDOM_STATE_OR_SEED_LIKE' = None,
+        seed: cirq.RANDOM_STATE_OR_SEED_LIKE = None,
     ) -> np.ndarray:
         if self._state is not None:
             return self._state.sample(self.get_axes(qubits), repetitions, seed)
@@ -166,7 +157,7 @@ class SimulationState(SimulationStateBase, Generic[TState], metaclass=abc.ABCMet
         """Creates a final merged state."""
         return self
 
-    def add_qubits(self: Self, qubits: Sequence['cirq.Qid']) -> Self:
+    def add_qubits(self: Self, qubits: Sequence[cirq.Qid]) -> Self:
         """Add `qubits` in the `|0>` state to a new state space and take the kron product.
 
         Args:
@@ -181,7 +172,7 @@ class SimulationState(SimulationStateBase, Generic[TState], metaclass=abc.ABCMet
             return self
         return NotImplemented
 
-    def remove_qubits(self: Self, qubits: Sequence['cirq.Qid']) -> Self:
+    def remove_qubits(self: Self, qubits: Sequence[cirq.Qid]) -> Self:
         """Remove `qubits` from the state space.
 
         The qubits to be removed should be untangled from rest of the system and in the |0> state.
@@ -206,8 +197,8 @@ class SimulationState(SimulationStateBase, Generic[TState], metaclass=abc.ABCMet
         return args
 
     def factor(
-        self, qubits: Sequence['cirq.Qid'], *, validate=True, atol=1e-07, inplace=False
-    ) -> Tuple[Self, Self]:
+        self, qubits: Sequence[cirq.Qid], *, validate=True, atol=1e-07, inplace=False
+    ) -> tuple[Self, Self]:
         """Splits two state spaces after a measurement or reset."""
         extracted = copy.copy(self)
         remainder = self if inplace else copy.copy(self)
@@ -223,7 +214,7 @@ class SimulationState(SimulationStateBase, Generic[TState], metaclass=abc.ABCMet
         """Subclasses that allow factorization should override this."""
         return self._state.supports_factor if self._state is not None else False
 
-    def transpose_to_qubit_order(self, qubits: Sequence['cirq.Qid'], *, inplace=False) -> Self:
+    def transpose_to_qubit_order(self, qubits: Sequence[cirq.Qid], *, inplace=False) -> Self:
         """Physically reindexes the state by the new basis.
 
         Args:
@@ -245,10 +236,10 @@ class SimulationState(SimulationStateBase, Generic[TState], metaclass=abc.ABCMet
         return args
 
     @property
-    def qubits(self) -> Tuple['cirq.Qid', ...]:
+    def qubits(self) -> tuple[cirq.Qid, ...]:
         return self._qubits
 
-    def swap(self, q1: 'cirq.Qid', q2: 'cirq.Qid', *, inplace=False):
+    def swap(self, q1: cirq.Qid, q2: cirq.Qid, *, inplace=False):
         """Swaps two qubits.
 
         This only affects the index, and does not modify the underlying
@@ -279,7 +270,7 @@ class SimulationState(SimulationStateBase, Generic[TState], metaclass=abc.ABCMet
         args._set_qubits(qubits)
         return args
 
-    def rename(self, q1: 'cirq.Qid', q2: 'cirq.Qid', *, inplace=False):
+    def rename(self, q1: cirq.Qid, q2: cirq.Qid, *, inplace=False):
         """Renames `q1` to `q2`.
 
         Args:
@@ -306,7 +297,7 @@ class SimulationState(SimulationStateBase, Generic[TState], metaclass=abc.ABCMet
         args._set_qubits(qubits)
         return args
 
-    def __getitem__(self, item: Optional['cirq.Qid']) -> Self:
+    def __getitem__(self, item: cirq.Qid | None) -> Self:
         if item not in self.qubit_map:
             raise IndexError(f'{item} not in {self.qubits}')
         return self
@@ -314,7 +305,7 @@ class SimulationState(SimulationStateBase, Generic[TState], metaclass=abc.ABCMet
     def __len__(self) -> int:
         return len(self.qubits)
 
-    def __iter__(self) -> Iterator[Optional['cirq.Qid']]:
+    def __iter__(self) -> Iterator[cirq.Qid | None]:
         return iter(self.qubits)
 
     @property
@@ -323,7 +314,7 @@ class SimulationState(SimulationStateBase, Generic[TState], metaclass=abc.ABCMet
 
 
 def strat_act_on_from_apply_decompose(
-    val: Any, args: 'cirq.SimulationState', qubits: Sequence['cirq.Qid']
+    val: Any, args: cirq.SimulationState, qubits: Sequence[cirq.Qid]
 ) -> bool:
     if isinstance(val, ops.Gate):
         decomposed = protocols.decompose_once_with_qubits(val, qubits, flatten=False, default=None)
@@ -331,7 +322,7 @@ def strat_act_on_from_apply_decompose(
         decomposed = protocols.decompose_once(val, flatten=False, default=None)
     if decomposed is None:
         return NotImplemented
-    all_ancilla: Set['cirq.Qid'] = set()
+    all_ancilla: set[cirq.Qid] = set()
     for operation in ops.flatten_to_ops(decomposed):
         curr_ancilla = tuple(q for q in operation.qubits if q not in args.qubits)
         args = args.add_qubits(curr_ancilla)

cirq/sim/simulation_state_base.py

@@ -11,30 +11,22 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
+
 """An interface for quantum states as targets for operations."""
+
+from __future__ import annotations
+
 import abc
 from types import NotImplementedType
-from typing import (
-    Any,
-    Dict,
-    Generic,
-    Iterator,
-    List,
-    Mapping,
-    Optional,
-    Sequence,
-    Tuple,
-    TYPE_CHECKING,
-    TypeVar,
-    Union,
-)
-
-import numpy as np
+from typing import Any, Generic, Iterator, Mapping, Sequence, TYPE_CHECKING, TypeVar
+
 from typing_extensions import Self
 
 from cirq import protocols, value
 
 if TYPE_CHECKING:
+    import numpy as np
+
     import cirq
 
 
@@ -45,10 +37,7 @@ class SimulationStateBase(Generic[TSimulationState], metaclass=abc.ABCMeta):
     """An interface for quantum states as targets for operations."""
 
     def __init__(
-        self,
-        *,
-        qubits: Sequence['cirq.Qid'],
-        classical_data: Optional['cirq.ClassicalDataStore'] = None,
+        self, *, qubits: Sequence[cirq.Qid], classical_data: cirq.ClassicalDataStore | None = None
     ):
         """Initializes the class.
 
@@ -61,19 +50,19 @@ class SimulationStateBase(Generic[TSimulationState], metaclass=abc.ABCMeta):
         self._classical_data = classical_data or value.ClassicalDataDictionaryStore()
 
     @property
-    def qubits(self) -> Tuple['cirq.Qid', ...]:
+    def qubits(self) -> tuple[cirq.Qid, ...]:
         return self._qubits
 
     @property
-    def qubit_map(self) -> Mapping['cirq.Qid', int]:
+    def qubit_map(self) -> Mapping[cirq.Qid, int]:
         return self._qubit_map
 
-    def _set_qubits(self, qubits: Sequence['cirq.Qid']):
+    def _set_qubits(self, qubits: Sequence[cirq.Qid]):
         self._qubits = tuple(qubits)
         self._qubit_map = {q: i for i, q in enumerate(self.qubits)}
 
     @property
-    def classical_data(self) -> 'cirq.ClassicalDataStoreReader':
+    def classical_data(self) -> cirq.ClassicalDataStoreReader:
         return self._classical_data
 
     @abc.abstractmethod
@@ -82,8 +71,8 @@ class SimulationStateBase(Generic[TSimulationState], metaclass=abc.ABCMeta):
 
     @abc.abstractmethod
     def _act_on_fallback_(
-        self, action: Any, qubits: Sequence['cirq.Qid'], allow_decompose: bool = True
-    ) -> Union[bool, NotImplementedType]:
+        self, action: Any, qubits: Sequence[cirq.Qid], allow_decompose: bool = True
+    ) -> bool | NotImplementedType:
         """Handles the act_on protocol fallback implementation.
 
         Args:
@@ -94,7 +83,7 @@ class SimulationStateBase(Generic[TSimulationState], metaclass=abc.ABCMeta):
         Returns:
             True if the fallback applies, else NotImplemented."""
 
-    def apply_operation(self, op: 'cirq.Operation'):
+    def apply_operation(self, op: cirq.Operation):
         protocols.act_on(op, self)
 
     @abc.abstractmethod
@@ -111,21 +100,21 @@ class SimulationStateBase(Generic[TSimulationState], metaclass=abc.ABCMeta):
         """
 
     @property
-    def log_of_measurement_results(self) -> Dict[str, List[int]]:
+    def log_of_measurement_results(self) -> dict[str, list[int]]:
         """Gets the log of measurement results."""
         return {str(k): list(self.classical_data.get_digits(k)) for k in self.classical_data.keys()}
 
     @abc.abstractmethod
     def sample(
         self,
-        qubits: List['cirq.Qid'],
+        qubits: list[cirq.Qid],
         repetitions: int = 1,
-        seed: 'cirq.RANDOM_STATE_OR_SEED_LIKE' = None,
+        seed: cirq.RANDOM_STATE_OR_SEED_LIKE = None,
     ) -> np.ndarray:
         """Samples the state value."""
 
     @abc.abstractmethod
-    def __getitem__(self, item: Optional['cirq.Qid']) -> TSimulationState:
+    def __getitem__(self, item: cirq.Qid | None) -> TSimulationState:
         """Gets the item associated with the qubit."""
 
     @abc.abstractmethod
@@ -133,5 +122,5 @@ class SimulationStateBase(Generic[TSimulationState], metaclass=abc.ABCMeta):
         """Gets the number of items in the mapping."""
 
     @abc.abstractmethod
-    def __iter__(self) -> Iterator[Optional['cirq.Qid']]:
+    def __iter__(self) -> Iterator[cirq.Qid | None]:
         """Iterates the keys of the mapping."""

cirq/sim/simulation_state_test.py

@@ -12,6 +12,8 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+from __future__ import annotations
+
 from typing import Any, Sequence
 
 import numpy as np
@@ -38,7 +40,7 @@ class ExampleSimulationState(cirq.SimulationState):
         super().__init__(state=ExampleQuantumState(), qubits=qubits)
 
     def _act_on_fallback_(
-        self, action: Any, qubits: Sequence['cirq.Qid'], allow_decompose: bool = True
+        self, action: Any, qubits: Sequence[cirq.Qid], allow_decompose: bool = True
     ) -> bool:
         return True
 
@@ -72,20 +74,20 @@ class Composite(cirq.Gate):
         yield cirq.X(*qubits)
 
 
-def test_measurements():
+def test_measurements() -> None:
     args = ExampleSimulationState()
     args.measure([cirq.LineQubit(0)], "test", [False], {})
     assert args.log_of_measurement_results["test"] == [5]
 
 
-def test_decompose():
+def test_decompose() -> None:
     args = ExampleSimulationState()
     assert simulation_state.strat_act_on_from_apply_decompose(
         Composite(), args, [cirq.LineQubit(0)]
     )
 
 
-def test_decompose_for_gate_allocating_qubits_raises():
+def test_decompose_for_gate_allocating_qubits_raises() -> None:
     class Composite(cirq.testing.SingleQubitGate):
         def _decompose_(self, qubits):
             anc = cirq.NamedQubit("anc")
@@ -97,7 +99,7 @@ def test_decompose_for_gate_allocating_qubits_raises():
         simulation_state.strat_act_on_from_apply_decompose(Composite(), args, [cirq.LineQubit(0)])
 
 
-def test_mapping():
+def test_mapping() -> None:
     args = ExampleSimulationState()
     assert list(iter(args)) == cirq.LineQubit.range(2)
     r1 = args[cirq.LineQubit(0)]
@@ -106,7 +108,7 @@ def test_mapping():
         _ = args[cirq.LineQubit(2)]
 
 
-def test_swap_bad_dimensions():
+def test_swap_bad_dimensions() -> None:
     q0 = cirq.LineQubit(0)
     q1 = cirq.LineQid(1, 3)
     args = ExampleSimulationState()
@@ -114,7 +116,7 @@ def test_swap_bad_dimensions():
         args.swap(q0, q1)
 
 
-def test_rename_bad_dimensions():
+def test_rename_bad_dimensions() -> None:
     q0 = cirq.LineQubit(0)
     q1 = cirq.LineQid(1, 3)
     args = ExampleSimulationState()
@@ -122,7 +124,7 @@ def test_rename_bad_dimensions():
         args.rename(q0, q1)
 
 
-def test_transpose_qubits():
+def test_transpose_qubits() -> None:
     q0, q1, q2 = cirq.LineQubit.range(3)
     args = ExampleSimulationState()
     assert args.transpose_to_qubit_order((q1, q0)).qubits == (q1, q0)
@@ -132,14 +134,14 @@ def test_transpose_qubits():
         args.transpose_to_qubit_order((q0, q1, q1))
 
 
-def test_field_getters():
+def test_field_getters() -> None:
     args = ExampleSimulationState()
     assert args.prng is np.random
     assert args.qubit_map == {q: i for i, q in enumerate(cirq.LineQubit.range(2))}
 
 
 @pytest.mark.parametrize('exp', np.linspace(0, 2 * np.pi, 10))
-def test_delegating_gate_unitary(exp):
+def test_delegating_gate_unitary(exp) -> None:
     q = cirq.LineQubit(0)
 
     test_circuit = cirq.Circuit()
@@ -154,7 +156,7 @@ def test_delegating_gate_unitary(exp):
 
 
 @pytest.mark.parametrize('exp', np.linspace(0, 2 * np.pi, 10))
-def test_delegating_gate_channel(exp):
+def test_delegating_gate_channel(exp) -> None:
     q = cirq.LineQubit(0)
 
     test_circuit = cirq.Circuit()
@@ -169,7 +171,7 @@ def test_delegating_gate_channel(exp):
 
 
 @pytest.mark.parametrize('num_ancilla', [1, 2, 3])
-def test_phase_using_dirty_ancilla(num_ancilla: int):
+def test_phase_using_dirty_ancilla(num_ancilla: int) -> None:
     q = cirq.LineQubit(0)
     anc = cirq.NamedQubit.range(num_ancilla, prefix='anc')
 
@@ -184,7 +186,7 @@ def test_phase_using_dirty_ancilla(num_ancilla: int):
 
 @pytest.mark.parametrize('num_ancilla', [1, 2, 3])
 @pytest.mark.parametrize('theta', np.linspace(0, 2 * np.pi, 10))
-def test_phase_using_clean_ancilla(num_ancilla: int, theta: float):
+def test_phase_using_clean_ancilla(num_ancilla: int, theta: float) -> None:
     q = cirq.LineQubit(0)
     u = cirq.MatrixGate(cirq.testing.random_unitary(2))
     test_circuit = cirq.Circuit(

cirq/sim/simulation_utils.py

@@ -11,7 +11,9 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
-from typing import Sequence, Tuple
+from __future__ import annotations
+
+from typing import Any, Sequence
 
 import numpy as np
 
@@ -19,7 +21,7 @@ from cirq import linalg
 
 
 def state_probabilities_by_indices(
-    state_probability: np.ndarray, indices: Sequence[int], qid_shape: Tuple[int, ...]
+    state_probability: np.ndarray, indices: Sequence[int], qid_shape: tuple[int, ...]
 ) -> np.ndarray:
     """Returns the probabilities for a state/measurement on the given indices.
 
@@ -37,6 +39,7 @@ def state_probabilities_by_indices(
     Returns:
         State probabilities.
     """
+    probs: np.ndarray[tuple[int, ...], Any]
     probs = state_probability.reshape((-1,))
     not_measured = [i for i in range(len(qid_shape)) if i not in indices]
     if linalg.can_numpy_support_shape(qid_shape):