cirq-core 1.7.0.dev20250915211227__py3-none-any.whl → 1.7.0.dev20250917002151__py3-none-any.whl

cirq/sim/density_matrix_utils_test.py

@@ -216,7 +216,7 @@ def test_measure_density_matrix_partial_indices_all_orders() -> None:
             assert bits == [bool(1 & (x >> (2 - p))) for p in perm]
 
 
-def matrix_000_plus_010():
+def matrix_000_plus_010() -> np.ndarray:
     state = np.zeros(8, dtype=np.complex64)
     state[0] = 1 / np.sqrt(2)
     state[2] = 1j / np.sqrt(2)

cirq/sim/mux_test.py

@@ -26,7 +26,7 @@ import cirq
 import cirq.testing
 
 
-def test_sample():
+def test_sample() -> None:
     q = cirq.NamedQubit('q')
 
     with pytest.raises(ValueError, match="no measurements"):
@@ -49,7 +49,7 @@ def test_sample():
     assert results.histogram(key=q) == collections.Counter({0: 1})
 
 
-def test_sample_seed_unitary():
+def test_sample_seed_unitary() -> None:
     q = cirq.NamedQubit('q')
     circuit = cirq.Circuit(cirq.X(q) ** 0.2, cirq.measure(q))
     result = cirq.sample(circuit, repetitions=10, seed=1234)
@@ -60,7 +60,7 @@ def test_sample_seed_unitary():
     )
 
 
-def test_sample_seed_non_unitary():
+def test_sample_seed_non_unitary() -> None:
     q = cirq.NamedQubit('q')
     circuit = cirq.Circuit(cirq.depolarize(0.5).on(q), cirq.measure(q))
     result = cirq.sample(circuit, repetitions=10, seed=1234)
@@ -70,7 +70,7 @@ def test_sample_seed_non_unitary():
     )
 
 
-def test_sample_sweep():
+def test_sample_sweep() -> None:
     q = cirq.NamedQubit('q')
     c = cirq.Circuit(cirq.X(q), cirq.Y(q) ** sympy.Symbol('t'), cirq.measure(q))
 
@@ -102,7 +102,7 @@ def test_sample_sweep():
     assert results[1].histogram(key=q) == collections.Counter({0: 3})
 
 
-def test_sample_sweep_seed():
+def test_sample_sweep_seed() -> None:
     q = cirq.NamedQubit('q')
     circuit = cirq.Circuit(cirq.X(q) ** sympy.Symbol('t'), cirq.measure(q))
 
@@ -124,7 +124,7 @@ def test_sample_sweep_seed():
     assert np.all(results[2].measurements['q'] == [[True], [False]])
 
 
-def test_final_state_vector_different_program_types():
+def test_final_state_vector_different_program_types() -> None:
     a, b = cirq.LineQubit.range(2)
 
     np.testing.assert_allclose(cirq.final_state_vector(cirq.X), [0, 1], atol=1e-8)
@@ -140,7 +140,7 @@ def test_final_state_vector_different_program_types():
     )
 
 
-def test_final_state_vector_initial_state():
+def test_final_state_vector_initial_state() -> None:
     np.testing.assert_allclose(cirq.final_state_vector(cirq.X, initial_state=0), [0, 1], atol=1e-8)
 
     np.testing.assert_allclose(cirq.final_state_vector(cirq.X, initial_state=1), [1, 0], atol=1e-8)
@@ -152,7 +152,7 @@ def test_final_state_vector_initial_state():
     )
 
 
-def test_final_state_vector_dtype_insensitive_to_initial_state():
+def test_final_state_vector_dtype_insensitive_to_initial_state() -> None:
     assert cirq.final_state_vector(cirq.X).dtype == np.complex64
 
     assert cirq.final_state_vector(cirq.X, initial_state=0).dtype == np.complex64
@@ -181,7 +181,7 @@ def test_final_state_vector_dtype_insensitive_to_initial_state():
         )
 
 
-def test_final_state_vector_param_resolver():
+def test_final_state_vector_param_resolver() -> None:
     s = sympy.Symbol('s')
 
     with pytest.raises(ValueError, match='not unitary'):
@@ -192,7 +192,7 @@ def test_final_state_vector_param_resolver():
     )
 
 
-def test_final_state_vector_qubit_order():
+def test_final_state_vector_qubit_order() -> None:
     a, b = cirq.LineQubit.range(2)
 
     np.testing.assert_allclose(
@@ -206,7 +206,7 @@ def test_final_state_vector_qubit_order():
     )
 
 
-def test_final_state_vector_ignore_terminal_measurement():
+def test_final_state_vector_ignore_terminal_measurement() -> None:
     a, b = cirq.LineQubit.range(2)
 
     np.testing.assert_allclose(
@@ -226,7 +226,7 @@ def test_final_state_vector_ignore_terminal_measurement():
 
 
 @pytest.mark.parametrize('repetitions', (0, 1, 100))
-def test_repetitions(repetitions):
+def test_repetitions(repetitions) -> None:
     a = cirq.LineQubit(0)
     c = cirq.Circuit(cirq.H(a), cirq.measure(a, key='m'))
     r = cirq.sample(c, repetitions=repetitions)
@@ -235,7 +235,7 @@ def test_repetitions(repetitions):
     assert np.issubdtype(samples.dtype, np.integer)
 
 
-def test_final_density_matrix_different_program_types():
+def test_final_density_matrix_different_program_types() -> None:
     a, b = cirq.LineQubit.range(2)
 
     np.testing.assert_allclose(cirq.final_density_matrix(cirq.X), [[0, 0], [0, 1]], atol=1e-8)
@@ -244,12 +244,12 @@ def test_final_density_matrix_different_program_types():
 
     np.testing.assert_allclose(
         cirq.final_density_matrix(cirq.Circuit(ops)),
-        [[0.5, 0, 0, 0.5], [0, 0, 0, 0], [0, 0, 0, 0], [0.5, 0, 0, 0.5]],
+        np.asarray([[0.5, 0, 0, 0.5], [0, 0, 0, 0], [0, 0, 0, 0], [0.5, 0, 0, 0.5]]),
         atol=1e-8,
     )
 
 
-def test_final_density_matrix_initial_state():
+def test_final_density_matrix_initial_state() -> None:
     np.testing.assert_allclose(
         cirq.final_density_matrix(cirq.X, initial_state=0), [[0, 0], [0, 1]], atol=1e-8
     )
@@ -265,7 +265,7 @@ def test_final_density_matrix_initial_state():
     )
 
 
-def test_final_density_matrix_dtype_insensitive_to_initial_state():
+def test_final_density_matrix_dtype_insensitive_to_initial_state() -> None:
     assert cirq.final_density_matrix(cirq.X).dtype == np.complex64
 
     assert cirq.final_density_matrix(cirq.X, initial_state=0).dtype == np.complex64
@@ -296,7 +296,7 @@ def test_final_density_matrix_dtype_insensitive_to_initial_state():
         )
 
 
-def test_final_density_matrix_param_resolver():
+def test_final_density_matrix_param_resolver() -> None:
     s = sympy.Symbol('s')
 
     with pytest.raises(ValueError, match='not specified in parameter sweep'):
@@ -308,17 +308,17 @@ def test_final_density_matrix_param_resolver():
     )
 
 
-def test_final_density_matrix_qubit_order():
+def test_final_density_matrix_qubit_order() -> None:
     a, b = cirq.LineQubit.range(2)
 
     np.testing.assert_allclose(
         cirq.final_density_matrix([cirq.X(a), cirq.X(b) ** 0.5], qubit_order=[a, b]),
-        [[0, 0, 0, 0], [0, 0, 0, 0], [0, 0, 0.5, 0.5j], [0, 0, -0.5j, 0.5]],
+        np.asarray([[0, 0, 0, 0], [0, 0, 0, 0], [0, 0, 0.5, 0.5j], [0, 0, -0.5j, 0.5]]),
     )
 
     np.testing.assert_allclose(
         cirq.final_density_matrix([cirq.X(a), cirq.X(b) ** 0.5], qubit_order=[b, a]),
-        [[0, 0, 0, 0], [0, 0.5, 0, 0.5j], [0, 0, 0, 0], [0, -0.5j, 0, 0.5]],
+        np.asarray([[0, 0, 0, 0], [0, 0.5, 0, 0.5j], [0, 0, 0, 0], [0, -0.5j, 0, 0.5]]),
     )
 
     np.testing.assert_allclose(
@@ -331,7 +331,7 @@ def test_final_density_matrix_qubit_order():
     )
 
 
-def test_final_density_matrix_seed_with_dephasing():
+def test_final_density_matrix_seed_with_dephasing() -> None:
     a = cirq.LineQubit(0)
     np.testing.assert_allclose(
         cirq.final_density_matrix([cirq.X(a) ** 0.5, cirq.measure(a)], seed=123),
@@ -345,7 +345,7 @@ def test_final_density_matrix_seed_with_dephasing():
     )
 
 
-def test_final_density_matrix_seed_with_collapsing():
+def test_final_density_matrix_seed_with_collapsing() -> None:
     a = cirq.LineQubit(0)
     np.testing.assert_allclose(
         cirq.final_density_matrix(
@@ -363,7 +363,7 @@ def test_final_density_matrix_seed_with_collapsing():
     )
 
 
-def test_final_density_matrix_noise():
+def test_final_density_matrix_noise() -> None:
     a = cirq.LineQubit(0)
     np.testing.assert_allclose(
         cirq.final_density_matrix([cirq.H(a), cirq.Z(a), cirq.H(a), cirq.measure(a)]),
@@ -380,7 +380,7 @@ def test_final_density_matrix_noise():
     )
 
 
-def test_final_density_matrix_classical_control():
+def test_final_density_matrix_classical_control() -> None:
     q0, q1 = cirq.LineQubit.range(2)
     circuit = cirq.Circuit(
         cirq.H(q0),
@@ -394,7 +394,7 @@ def test_final_density_matrix_classical_control():
     )
 
 
-def test_ps_initial_state_wfn():
+def test_ps_initial_state_wfn() -> None:
     q0, q1 = cirq.LineQubit.range(2)
     s00 = cirq.KET_ZERO(q0) * cirq.KET_ZERO(q1)
     sp0 = cirq.KET_PLUS(q0) * cirq.KET_ZERO(q1)
@@ -410,7 +410,7 @@ def test_ps_initial_state_wfn():
     )
 
 
-def test_ps_initial_state_dmat():
+def test_ps_initial_state_dmat() -> None:
     q0, q1 = cirq.LineQubit.range(2)
     s00 = cirq.KET_ZERO(q0) * cirq.KET_ZERO(q1)
     sp0 = cirq.KET_PLUS(q0) * cirq.KET_ZERO(q1)

cirq/sim/simulation_product_state_test.py

@@ -14,29 +14,29 @@
 
 from __future__ import annotations
 
-from typing import Any, Sequence
+from typing import Any, Self, Sequence
 
 import cirq
 
 
 class EmptyQuantumState(cirq.QuantumStateRepresentation):
-    def copy(self, deep_copy_buffers=True):
+    def copy(self, deep_copy_buffers=True) -> Self:
         return self
 
-    def measure(self, axes, seed=None):
+    def measure(self, axes, seed=None) -> list[int]:
         return [0] * len(axes)
 
     @property
-    def supports_factor(self):
+    def supports_factor(self) -> bool:
         return True
 
-    def kron(self, other):
+    def kron(self, other) -> Self:
         return self
 
-    def factor(self, axes, *, validate=True, atol=1e-07):
+    def factor(self, axes, *, validate=True, atol=1e-07) -> tuple[Self, Self]:
         return self, self
 
-    def reindex(self, axes):
+    def reindex(self, axes) -> Self:
         return self
 
 

cirq/sim/simulation_state.py

@@ -73,7 +73,7 @@ class SimulationState(SimulationStateBase, Generic[TState], metaclass=abc.ABCMet
         key: str,
         invert_mask: Sequence[bool],
         confusion_map: dict[tuple[int, ...], np.ndarray],
-    ):
+    ) -> None:
         """Measures the qubits and records to `log_of_measurement_results`.
 
         Any bitmasks will be applied to the measurement record.
@@ -209,7 +209,7 @@ class SimulationState(SimulationStateBase, Generic[TState], metaclass=abc.ABCMet
         return extracted, remainder
 
     @property
-    def allows_factoring(self):
+    def allows_factoring(self) -> bool:
         """Subclasses that allow factorization should override this."""
         return self._state.supports_factor if self._state is not None else False
 
@@ -238,7 +238,7 @@ class SimulationState(SimulationStateBase, Generic[TState], metaclass=abc.ABCMet
     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) -> Self:
         """Swaps two qubits.
 
         This only affects the index, and does not modify the underlying
@@ -269,7 +269,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) -> Self:
         """Renames `q1` to `q2`.
 
         Args:

cirq/sim/simulation_state_base.py

@@ -81,7 +81,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) -> None:
         protocols.act_on(op, self)
 
     @abc.abstractmethod

cirq/sim/simulation_state_test.py

@@ -25,13 +25,13 @@ from cirq.testing import PhaseUsingCleanAncilla, PhaseUsingDirtyAncilla
 
 
 class ExampleQuantumState(cirq.QuantumStateRepresentation):
-    def copy(self, deep_copy_buffers=True):
-        pass
+    def copy(self, deep_copy_buffers=True) -> ExampleQuantumState:
+        raise NotImplementedError()
 
-    def measure(self, axes, seed=None):
+    def measure(self, axes, seed=None) -> list[int]:
         return [5, 3]
 
-    def reindex(self, axes):
+    def reindex(self, axes) -> ExampleQuantumState:
         return self
 
 
@@ -44,7 +44,7 @@ class ExampleSimulationState(cirq.SimulationState):
     ) -> bool:
         return True
 
-    def add_qubits(self, qubits):
+    def add_qubits(self, qubits) -> ExampleSimulationState:
         super().add_qubits(qubits)
         return self
 

cirq/sim/simulator.py

@@ -217,7 +217,7 @@ class SimulatesAmplitudes(metaclass=value.ABCMetaImplementAnyOneOf):
     def sample_from_amplitudes(
         self,
         circuit: cirq.AbstractCircuit,
-        param_resolver: cirq.ParamResolver,
+        param_resolver: cirq.ParamResolverOrSimilarType,
         seed: cirq.RANDOM_STATE_OR_SEED_LIKE,
         repetitions: int = 1,
         qubit_order: cirq.QubitOrderOrList = ops.QubitOrder.DEFAULT,
@@ -938,7 +938,7 @@ def _qubit_map_to_shape(qubit_map: Mapping[cirq.Qid, int]) -> tuple[int, ...]:
     return tuple(qid_shape)
 
 
-def check_all_resolved(circuit):
+def check_all_resolved(circuit) -> None:
     """Raises if the circuit contains unresolved symbols."""
     if protocols.is_parameterized(circuit):
         unresolved = [op for moment in circuit for op in moment if protocols.is_parameterized(op)]

cirq/sim/simulator_base_test.py

@@ -73,25 +73,25 @@ class CountingSimulationState(cirq.SimulationState[CountingState]):
         return True
 
     @property
-    def data(self):
+    def data(self) -> Any:
         return self._state.data
 
     @property
-    def gate_count(self):
+    def gate_count(self) -> int:
         return self._state.gate_count
 
     @property
-    def measurement_count(self):
+    def measurement_count(self) -> int:
         return self._state.measurement_count
 
     @property
-    def copy_count(self):
+    def copy_count(self) -> int:
         return self._state.copy_count
 
 
 class SplittableCountingSimulationState(CountingSimulationState):
     @property
-    def allows_factoring(self):
+    def allows_factoring(self) -> bool:
         return True
 
 
@@ -167,59 +167,64 @@ entangled_state_repr = np.array([[math.sqrt(0.5), 0], [0, math.sqrt(0.5)]])
 
 
 class TestOp(cirq.Operation):
-    def with_qubits(self, *new_qubits):
-        pass
+    def with_qubits(self, *new_qubits) -> cirq.Operation:
+        raise NotImplementedError()
 
     @property
-    def qubits(self):
-        return [q0]
+    def qubits(self) -> tuple[cirq.Qid, ...]:
+        return (q0,)
 
 
-def test_simulate_empty_circuit():
+def test_simulate_empty_circuit() -> None:
     sim = CountingSimulator()
     r = sim.simulate(cirq.Circuit())
+    assert isinstance(r._final_simulator_state, CountingSimulationState)
     assert r._final_simulator_state.gate_count == 0
     assert r._final_simulator_state.measurement_count == 0
     assert r._final_simulator_state.copy_count == 0
 
 
-def test_simulate_one_gate_circuit():
+def test_simulate_one_gate_circuit() -> None:
     sim = CountingSimulator()
     r = sim.simulate(cirq.Circuit(cirq.X(q0)))
+    assert isinstance(r._final_simulator_state, CountingSimulationState)
     assert r._final_simulator_state.gate_count == 1
     assert r._final_simulator_state.copy_count == 0
 
 
-def test_simulate_one_measurement_circuit():
+def test_simulate_one_measurement_circuit() -> None:
     sim = CountingSimulator()
     r = sim.simulate(cirq.Circuit(cirq.measure(q0)))
+    assert isinstance(r._final_simulator_state, CountingSimulationState)
     assert r._final_simulator_state.gate_count == 0
     assert r._final_simulator_state.measurement_count == 1
     assert r._final_simulator_state.copy_count == 0
 
 
-def test_empty_circuit_simulation_has_moment():
+def test_empty_circuit_simulation_has_moment() -> None:
     sim = CountingSimulator()
     steps = list(sim.simulate_moment_steps(cirq.Circuit()))
     assert len(steps) == 1
 
 
-def test_noise_applied():
+def test_noise_applied() -> None:
     sim = CountingSimulator(noise=cirq.X)
     r = sim.simulate(cirq.Circuit(cirq.X(q0)))
+    assert isinstance(r._final_simulator_state, CountingSimulationState)
     assert r._final_simulator_state.gate_count == 2
     assert r._final_simulator_state.copy_count == 0
 
 
-def test_noise_applied_measurement_gate():
+def test_noise_applied_measurement_gate() -> None:
     sim = CountingSimulator(noise=cirq.X)
     r = sim.simulate(cirq.Circuit(cirq.measure(q0)))
+    assert isinstance(r._final_simulator_state, CountingSimulationState)
     assert r._final_simulator_state.gate_count == 1
     assert r._final_simulator_state.measurement_count == 1
     assert r._final_simulator_state.copy_count == 0
 
 
-def test_parameterized_copies_all_but_last():
+def test_parameterized_copies_all_but_last() -> None:
     sim = CountingSimulator()
     n = 4
     rs = sim.simulate_sweep(
@@ -227,12 +232,13 @@ def test_parameterized_copies_all_but_last():
     )
     for i in range(n):
         r = rs[i]
+        assert isinstance(r._final_simulator_state, CountingSimulationState)
         assert r._final_simulator_state.gate_count == 1
         assert r._final_simulator_state.measurement_count == 0
         assert r._final_simulator_state.copy_count == 0 if i == n - 1 else 1
 
 
-def test_cannot_act():
+def test_cannot_act() -> None:
     class BadOp(TestOp):
         def _act_on_(self, sim_state):
             raise TypeError()
@@ -242,25 +248,25 @@ def test_cannot_act():
         sim.simulate(cirq.Circuit(BadOp()))
 
 
-def test_run_one_gate_circuit():
+def test_run_one_gate_circuit() -> None:
     sim = CountingSimulator()
     r = sim.run(cirq.Circuit(cirq.X(q0), cirq.measure(q0)), repetitions=2)
     assert np.allclose(r.measurements['q(0)'], [[1], [1]])
 
 
-def test_run_one_gate_circuit_noise():
+def test_run_one_gate_circuit_noise() -> None:
     sim = CountingSimulator(noise=cirq.X)
     r = sim.run(cirq.Circuit(cirq.X(q0), cirq.measure(q0)), repetitions=2)
     assert np.allclose(r.measurements['q(0)'], [[2], [2]])
 
 
-def test_run_non_unitary_circuit():
+def test_run_non_unitary_circuit() -> None:
     sim = CountingSimulator()
     r = sim.run(cirq.Circuit(cirq.phase_damp(1).on(q0), cirq.measure(q0)), repetitions=2)
     assert np.allclose(r.measurements['q(0)'], [[1], [1]])
 
 
-def test_run_non_unitary_circuit_non_unitary_state():
+def test_run_non_unitary_circuit_non_unitary_state() -> None:
     class DensityCountingSimulator(CountingSimulator):
         def _can_be_in_run_prefix(self, val):
             return not cirq.is_measurement(val)
@@ -270,15 +276,16 @@ def test_run_non_unitary_circuit_non_unitary_state():
     assert np.allclose(r.measurements['q(0)'], [[1], [1]])
 
 
-def test_run_non_terminal_measurement():
+def test_run_non_terminal_measurement() -> None:
     sim = CountingSimulator()
     r = sim.run(cirq.Circuit(cirq.X(q0), cirq.measure(q0), cirq.X(q0)), repetitions=2)
     assert np.allclose(r.measurements['q(0)'], [[1], [1]])
 
 
-def test_integer_initial_state_is_split():
+def test_integer_initial_state_is_split() -> None:
     sim = SplittableCountingSimulator()
     state = sim._create_simulation_state(2, (q0, q1))
+    assert isinstance(state, cirq.SimulationProductState)
     assert len(set(state.values())) == 3
     assert state[q0] is not state[q1]
     assert state[q0].data == 1
@@ -286,7 +293,7 @@ def test_integer_initial_state_is_split():
     assert state[None].data == 0
 
 
-def test_integer_initial_state_is_not_split_if_disabled():
+def test_integer_initial_state_is_not_split_if_disabled() -> None:
     sim = SplittableCountingSimulator(split_untangled_states=False)
     state = sim._create_simulation_state(2, (q0, q1))
     assert isinstance(state, SplittableCountingSimulationState)
@@ -294,7 +301,7 @@ def test_integer_initial_state_is_not_split_if_disabled():
     assert state.data == 2
 
 
-def test_integer_initial_state_is_not_split_if_impossible():
+def test_integer_initial_state_is_not_split_if_impossible() -> None:
     sim = CountingSimulator()
     state = sim._create_simulation_state(2, (q0, q1))
     assert isinstance(state, CountingSimulationState)
@@ -303,18 +310,20 @@ def test_integer_initial_state_is_not_split_if_impossible():
     assert state.data == 2
 
 
-def test_non_integer_initial_state_is_not_split():
+def test_non_integer_initial_state_is_not_split() -> None:
     sim = SplittableCountingSimulator()
     state = sim._create_simulation_state(entangled_state_repr, (q0, q1))
+    assert isinstance(state, cirq.SimulationProductState)
     assert len(set(state.values())) == 2
     assert (state[q0].data == entangled_state_repr).all()
     assert state[q1] is state[q0]
     assert state[None].data == 0
 
 
-def test_entanglement_causes_join():
+def test_entanglement_causes_join() -> None:
     sim = SplittableCountingSimulator()
     state = sim._create_simulation_state(2, (q0, q1))
+    assert isinstance(state, cirq.SimulationProductState)
     assert len(set(state.values())) == 3
     state.apply_operation(cirq.CNOT(q0, q1))
     assert len(set(state.values())) == 2
@@ -322,9 +331,10 @@ def test_entanglement_causes_join():
     assert state[None] is not state[q0]
 
 
-def test_measurement_causes_split():
+def test_measurement_causes_split() -> None:
     sim = SplittableCountingSimulator()
     state = sim._create_simulation_state(entangled_state_repr, (q0, q1))
+    assert isinstance(state, cirq.SimulationProductState)
     assert len(set(state.values())) == 2
     state.apply_operation(cirq.measure(q0))
     assert len(set(state.values())) == 3
@@ -332,7 +342,7 @@ def test_measurement_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:
     sim = SplittableCountingSimulator(split_untangled_states=False)
     state = sim._create_simulation_state(2, (q0, q1))
     assert isinstance(state, SplittableCountingSimulationState)
@@ -341,7 +351,7 @@ def test_measurement_does_not_split_if_disabled():
     assert state[q0] is state[q1]
 
 
-def test_measurement_does_not_split_if_impossible():
+def test_measurement_does_not_split_if_impossible() -> None:
     sim = CountingSimulator()
     state = sim._create_simulation_state(2, (q0, q1))
     assert isinstance(state, CountingSimulationState)
@@ -352,7 +362,7 @@ def test_measurement_does_not_split_if_impossible():
     assert state[q0] is state[q1]
 
 
-def test_reorder_succeeds():
+def test_reorder_succeeds() -> None:
     sim = SplittableCountingSimulator()
     state = sim._create_simulation_state(entangled_state_repr, (q0, q1))
     reordered = state[q0].transpose_to_qubit_order([q1, q0])
@@ -360,7 +370,7 @@ def test_reorder_succeeds():
 
 
 @pytest.mark.parametrize('split', [True, False])
-def test_sim_state_instance_unchanged_during_normal_sim(split: bool):
+def test_sim_state_instance_unchanged_during_normal_sim(split: bool) -> None:
     sim = SplittableCountingSimulator(split_untangled_states=split)
     state = sim._create_simulation_state(0, (q0, q1))
     circuit = cirq.Circuit(cirq.H(q0), cirq.CNOT(q0, q1), cirq.reset(q1))
@@ -369,7 +379,7 @@ def test_sim_state_instance_unchanged_during_normal_sim(split: bool):
         assert (step._merged_sim_state is not state) == split
 
 
-def test_measurements_retained_in_step_results():
+def test_measurements_retained_in_step_results() -> None:
     sim = SplittableCountingSimulator()
     circuit = cirq.Circuit(
         cirq.measure(q0, key='a'), cirq.measure(q0, key='b'), cirq.measure(q0, key='c')
@@ -381,7 +391,7 @@ def test_measurements_retained_in_step_results():
     assert not any(iterator)
 
 
-def test_sweep_unparameterized_prefix_not_repeated_iff_unitary():
+def test_sweep_unparameterized_prefix_not_repeated_iff_unitary() -> None:
     q = cirq.LineQubit(0)
 
     class TestOp(cirq.Operation):
@@ -410,6 +420,8 @@ def test_sweep_unparameterized_prefix_not_repeated_iff_unitary():
     op2 = TestOp(has_unitary=True)
     circuit = cirq.Circuit(op1, cirq.XPowGate(exponent=sympy.Symbol('a'))(q), op2)
     rs = simulator.simulate_sweep(program=circuit, params=params)
+    assert isinstance(rs[0]._final_simulator_state, CountingSimulationState)
+    assert isinstance(rs[1]._final_simulator_state, CountingSimulationState)
     assert rs[0]._final_simulator_state.copy_count == 1
     assert rs[1]._final_simulator_state.copy_count == 0
     assert op1.count == 1
@@ -419,13 +431,15 @@ def test_sweep_unparameterized_prefix_not_repeated_iff_unitary():
     op2 = TestOp(has_unitary=False)
     circuit = cirq.Circuit(op1, cirq.XPowGate(exponent=sympy.Symbol('a'))(q), op2)
     rs = simulator.simulate_sweep(program=circuit, params=params)
+    assert isinstance(rs[0]._final_simulator_state, CountingSimulationState)
+    assert isinstance(rs[1]._final_simulator_state, CountingSimulationState)
     assert rs[0]._final_simulator_state.copy_count == 1
     assert rs[1]._final_simulator_state.copy_count == 0
     assert op1.count == 2
     assert op2.count == 2
 
 
-def test_inhomogeneous_measurement_count_padding():
+def test_inhomogeneous_measurement_count_padding() -> None:
     q = cirq.LineQubit(0)
     key = cirq.MeasurementKey('m')
     sim = cirq.Simulator()