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

cirq/linalg/transformations_test.py

@@ -12,6 +12,8 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+from __future__ import annotations
+
 import numpy as np
 import pytest
 
@@ -19,8 +21,10 @@ import cirq
 import cirq.testing
 from cirq import linalg
 
+_DEFAULT_ARRAY = np.asarray([1])
+
 
-def test_reflection_matrix_pow_consistent_results():
+def test_reflection_matrix_pow_consistent_results() -> None:
     x = np.array([[0, 1], [1, 0]])
     sqrt_x = cirq.reflection_matrix_pow(x, 0.5)
     np.testing.assert_allclose(np.dot(sqrt_x, sqrt_x), x, atol=1e-10)
@@ -40,7 +44,7 @@ def test_reflection_matrix_pow_consistent_results():
     np.testing.assert_allclose(np.dot(sqrt_yh, sqrt_yh), yh, atol=1e-10)
 
 
-def test_reflection_matrix_sign_preference_under_perturbation():
+def test_reflection_matrix_sign_preference_under_perturbation() -> None:
     x = np.array([[0, 1], [1, 0]])
     sqrt_x = np.array([[1, -1j], [-1j, 1]]) * (1 + 1j) / 2
     np.testing.assert_allclose(cirq.reflection_matrix_pow(x, 0.5), sqrt_x, atol=1e-8)
@@ -55,7 +59,7 @@ def test_reflection_matrix_sign_preference_under_perturbation():
         np.testing.assert_allclose(sqrt_px, expected_sqrt_px, atol=1e-10)
 
 
-def test_match_global_phase():
+def test_match_global_phase() -> None:
     a = np.array([[5, 4], [3, -2]])
     b = np.array([[0.000001, 0], [0, 1j]])
     c, d = cirq.match_global_phase(a, b)
@@ -63,7 +67,7 @@ def test_match_global_phase():
     np.testing.assert_allclose(d, b * -1j, atol=1e-10)
 
 
-def test_match_global_phase_incompatible_shape():
+def test_match_global_phase_incompatible_shape() -> None:
     a = np.array([1])
     b = np.array([1, 2])
     c, d = cirq.match_global_phase(a, b)
@@ -85,7 +89,7 @@ def test_match_global_phase_incompatible_shape():
     assert np.allclose(d, b)
 
 
-def test_match_global_phase_zeros():
+def test_match_global_phase_zeros() -> None:
     z = np.array([[0, 0], [0, 0]])
     b = np.array([[1, 1], [1, 1]])
 
@@ -102,7 +106,7 @@ def test_match_global_phase_zeros():
     np.testing.assert_allclose(z, z4, atol=1e-10)
 
 
-def test_match_global_no_float_error_when_axis_aligned():
+def test_match_global_no_float_error_when_axis_aligned() -> None:
     a = np.array([[1, 1.1], [-1.3, np.pi]])
     a2, _ = cirq.match_global_phase(a, a)
     a3, _ = cirq.match_global_phase(a * 1j, a * 1j)
@@ -115,7 +119,7 @@ def test_match_global_no_float_error_when_axis_aligned():
     assert np.all(a5 == a)
 
 
-def test_targeted_left_multiply_matches_kron_then_dot():
+def test_targeted_left_multiply_matches_kron_then_dot() -> None:
     t = np.array([1, 2, 3, 4, 5, 6, 7, 8])
     m = np.array([[2, 3], [5, 7]])
     i = np.eye(2)
@@ -159,7 +163,7 @@ def test_targeted_left_multiply_matches_kron_then_dot():
         cirq.targeted_left_multiply(left_matrix=m, right_target=common, out=m, target_axes=[2])
 
 
-def test_targeted_left_multiply_reorders_matrices():
+def test_targeted_left_multiply_reorders_matrices() -> None:
     t = np.eye(4).reshape((2, 2, 2, 2))
     m = np.array([1, 0, 0, 0, 0, 1, 0, 0, 0, 0, 0, 1, 0, 0, 1, 0]).reshape((2, 2, 2, 2))
 
@@ -174,7 +178,7 @@ def test_targeted_left_multiply_reorders_matrices():
     )
 
 
-def test_targeted_left_multiply_out():
+def test_targeted_left_multiply_out() -> None:
     left = np.array([[2, 3], [5, 7]])
     right = np.array([1, -1])
     out = np.zeros(2)
@@ -186,7 +190,7 @@ def test_targeted_left_multiply_out():
     np.testing.assert_allclose(result, np.array([-1, -2]), atol=1e-8)
 
 
-def test_targeted_conjugate_simple():
+def test_targeted_conjugate_simple() -> None:
     a = np.array([[0, 1j], [0, 0]])
     # yapf: disable
     b = np.reshape(
@@ -210,7 +214,7 @@ def test_targeted_conjugate_simple():
     np.testing.assert_almost_equal(result, expected)
 
 
-def test_targeted_conjugate():
+def test_targeted_conjugate() -> None:
     a = np.reshape([0, 1, 2j, 3j], (2, 2))
     b = np.reshape(np.arange(16), (2,) * 4)
     result = cirq.targeted_conjugate_about(a, b, [0])
@@ -222,7 +226,7 @@ def test_targeted_conjugate():
     np.testing.assert_almost_equal(result, expected)
 
 
-def test_targeted_conjugate_multiple_indices():
+def test_targeted_conjugate_multiple_indices() -> None:
     a = np.reshape(np.arange(16) + 1j, (2, 2, 2, 2))
     b = np.reshape(np.arange(16), (2,) * 4)
     result = cirq.targeted_conjugate_about(a, b, [0, 1])
@@ -230,7 +234,7 @@ def test_targeted_conjugate_multiple_indices():
     np.testing.assert_almost_equal(result, expected)
 
 
-def test_targeted_conjugate_multiple_indices_flip_order():
+def test_targeted_conjugate_multiple_indices_flip_order() -> None:
     a = np.reshape(np.arange(16) + 1j, (2, 2, 2, 2))
     b = np.reshape(np.arange(16), (2,) * 4)
     result = cirq.targeted_conjugate_about(a, b, [1, 0], [3, 2])
@@ -238,7 +242,7 @@ def test_targeted_conjugate_multiple_indices_flip_order():
     np.testing.assert_almost_equal(result, expected)
 
 
-def test_targeted_conjugate_out():
+def test_targeted_conjugate_out() -> None:
     a = np.reshape([0, 1, 2j, 3j], (2, 2))
     b = np.reshape(np.arange(16), (2,) * 4)
     buffer = np.empty((2,) * 4, dtype=a.dtype)
@@ -249,7 +253,7 @@ def test_targeted_conjugate_out():
     np.testing.assert_almost_equal(result, expected)
 
 
-def test_apply_matrix_to_slices():
+def test_apply_matrix_to_slices() -> None:
     # Output is input.
     with pytest.raises(ValueError, match='out'):
         target = np.eye(2)
@@ -301,7 +305,7 @@ def test_apply_matrix_to_slices():
     np.testing.assert_allclose(actual, np.array([1, 13, 31, 4]))
 
 
-def test_partial_trace():
+def test_partial_trace() -> None:
     a = np.reshape(np.arange(4), (2, 2))
     b = np.reshape(np.arange(9) + 4, (3, 3))
     c = np.reshape(np.arange(16) + 13, (4, 4))
@@ -335,21 +339,21 @@ def test_partial_trace():
     )
 
 
-def test_partial_trace_non_kron():
+def test_partial_trace_non_kron() -> None:
     tensor = np.zeros((2, 2, 2, 2))
     tensor[0, 0, 0, 0] = 1
     tensor[1, 1, 1, 1] = 4
     np.testing.assert_almost_equal(cirq.partial_trace(tensor, [0]), np.array([[1, 0], [0, 4]]))
 
 
-def test_partial_trace_invalid_inputs():
+def test_partial_trace_invalid_inputs() -> None:
     with pytest.raises(ValueError, match='2, 3, 2, 2'):
         cirq.partial_trace(np.reshape(np.arange(2 * 3 * 2 * 2), (2, 3, 2, 2)), [1])
     with pytest.raises(ValueError, match='2'):
         cirq.partial_trace(np.reshape(np.arange(2 * 2 * 2 * 2), (2,) * 4), [2])
 
 
-def test_sub_state_vector():
+def test_sub_state_vector() -> None:
     a = np.arange(4) / np.linalg.norm(np.arange(4))
     b = (np.arange(8) + 3) / np.linalg.norm(np.arange(8) + 3)
     c = (np.arange(16) + 1) / np.linalg.norm(np.arange(16) + 1)
@@ -382,40 +386,66 @@ def test_sub_state_vector():
     )
 
     # Reject factoring for very tight tolerance.
-    assert cirq.sub_state_vector(state, [0, 1], default=None, atol=1e-16) is None
-    assert cirq.sub_state_vector(state, [2, 3, 4], default=None, atol=1e-16) is None
-    assert cirq.sub_state_vector(state, [5, 6, 7, 8], default=None, atol=1e-16) is None
+    assert (
+        cirq.sub_state_vector(state, [0, 1], default=_DEFAULT_ARRAY, atol=1e-16) is _DEFAULT_ARRAY
+    )
+    assert (
+        cirq.sub_state_vector(state, [2, 3, 4], default=_DEFAULT_ARRAY, atol=1e-16)
+        is _DEFAULT_ARRAY
+    )
+    assert (
+        cirq.sub_state_vector(state, [5, 6, 7, 8], default=_DEFAULT_ARRAY, atol=1e-16)
+        is _DEFAULT_ARRAY
+    )
 
     # Permit invalid factoring for loose tolerance.
     for q1 in range(9):
-        assert cirq.sub_state_vector(state, [q1], default=None, atol=1) is not None
+        assert (
+            cirq.sub_state_vector(state, [q1], default=_DEFAULT_ARRAY, atol=1) is not _DEFAULT_ARRAY
+        )
 
 
-def test_sub_state_vector_bad_subset():
+def test_sub_state_vector_bad_subset() -> None:
     a = cirq.testing.random_superposition(4)
     b = cirq.testing.random_superposition(8)
     state = np.kron(a, b).reshape((2, 2, 2, 2, 2))
 
     for q1 in range(5):
-        assert cirq.sub_state_vector(state, [q1], default=None, atol=1e-8) is None
+        assert (
+            cirq.sub_state_vector(state, [q1], default=_DEFAULT_ARRAY, atol=1e-8) is _DEFAULT_ARRAY
+        )
     for q1 in range(2):
         for q2 in range(2, 5):
-            assert cirq.sub_state_vector(state, [q1, q2], default=None, atol=1e-8) is None
+            assert (
+                cirq.sub_state_vector(state, [q1, q2], default=_DEFAULT_ARRAY, atol=1e-8)
+                is _DEFAULT_ARRAY
+            )
     for q3 in range(2, 5):
-        assert cirq.sub_state_vector(state, [0, 1, q3], default=None, atol=1e-8) is None
+        assert (
+            cirq.sub_state_vector(state, [0, 1, q3], default=_DEFAULT_ARRAY, atol=1e-8)
+            is _DEFAULT_ARRAY
+        )
     for q4 in range(2):
-        assert cirq.sub_state_vector(state, [2, 3, 4, q4], default=None, atol=1e-8) is None
+        assert (
+            cirq.sub_state_vector(state, [2, 3, 4, q4], default=_DEFAULT_ARRAY, atol=1e-8)
+            is _DEFAULT_ARRAY
+        )
 
 
-def test_sub_state_vector_non_kron():
+def test_sub_state_vector_non_kron() -> None:
     a = np.array([1, 0, 0, 0, 0, 0, 0, 1]) / np.sqrt(2)
     b = np.array([1, 1]) / np.sqrt(2)
     state = np.kron(a, b).reshape((2, 2, 2, 2))
 
     for q1 in [0, 1, 2]:
-        assert cirq.sub_state_vector(state, [q1], default=None, atol=1e-8) is None
+        assert (
+            cirq.sub_state_vector(state, [q1], default=_DEFAULT_ARRAY, atol=1e-8) is _DEFAULT_ARRAY
+        )
     for q1 in [0, 1, 2]:
-        assert cirq.sub_state_vector(state, [q1, 3], default=None, atol=1e-8) is None
+        assert (
+            cirq.sub_state_vector(state, [q1, 3], default=_DEFAULT_ARRAY, atol=1e-8)
+            is _DEFAULT_ARRAY
+        )
 
     with pytest.raises(ValueError, match='factored'):
         _ = cirq.sub_state_vector(a, [0], atol=1e-8)
@@ -423,7 +453,7 @@ def test_sub_state_vector_non_kron():
     assert cirq.equal_up_to_global_phase(cirq.sub_state_vector(state, [3]), b, atol=1e-8)
 
 
-def test_sub_state_vector_invalid_inputs():
+def test_sub_state_vector_invalid_inputs() -> None:
 
     # State cannot be expressed as a separable pure state.
     with pytest.raises(ValueError, match='7'):
@@ -450,7 +480,7 @@ def test_sub_state_vector_invalid_inputs():
         cirq.sub_state_vector(state, [0, 1, 2], atol=1e-8)
 
 
-def test_partial_trace_of_state_vector_as_mixture_invalid_input():
+def test_partial_trace_of_state_vector_as_mixture_invalid_input() -> None:
 
     with pytest.raises(ValueError, match='7'):
         cirq.partial_trace_of_state_vector_as_mixture(np.arange(7), [1, 2], atol=1e-8)
@@ -469,7 +499,7 @@ def test_partial_trace_of_state_vector_as_mixture_invalid_input():
         cirq.partial_trace_of_state_vector_as_mixture(state, [0, 1, 2], atol=1e-8)
 
 
-def mixtures_equal(m1, m2, atol=1e-7):
+def mixtures_equal(m1, m2, atol=1e-7) -> bool:
     for (p1, v1), (p2, v2) in zip(m1, m2):
         if not (
             cirq.approx_eq(p1, p2, atol=atol) and cirq.equal_up_to_global_phase(v1, v2, atol=atol)
@@ -478,7 +508,7 @@ def mixtures_equal(m1, m2, atol=1e-7):
     return True
 
 
-def test_partial_trace_of_state_vector_as_mixture_pure_result():
+def test_partial_trace_of_state_vector_as_mixture_pure_result() -> None:
     a = cirq.testing.random_superposition(4)
     b = cirq.testing.random_superposition(8)
     c = cirq.testing.random_superposition(16)
@@ -532,7 +562,7 @@ def test_partial_trace_of_state_vector_as_mixture_pure_result():
     )
 
 
-def test_partial_trace_of_state_vector_as_mixture_pure_result_qudits():
+def test_partial_trace_of_state_vector_as_mixture_pure_result_qudits() -> None:
     a = cirq.testing.random_superposition(2)
     b = cirq.testing.random_superposition(3)
     c = cirq.testing.random_superposition(4)
@@ -561,7 +591,7 @@ def test_partial_trace_of_state_vector_as_mixture_pure_result_qudits():
     )
 
 
-def test_partial_trace_of_state_vector_as_mixture_mixed_result():
+def test_partial_trace_of_state_vector_as_mixture_mixed_result() -> None:
     state = np.array([1, 0, 0, 1]) / np.sqrt(2)
     truth = ((0.5, np.array([1, 0])), (0.5, np.array([0, 1])))
     for q1 in [0, 1]:
@@ -589,7 +619,7 @@ def test_partial_trace_of_state_vector_as_mixture_mixed_result():
         assert mixtures_equal(mixture, truth)
 
 
-def test_partial_trace_of_state_vector_as_mixture_mixed_result_qudits():
+def test_partial_trace_of_state_vector_as_mixture_mixed_result_qudits() -> None:
     state = np.array([[1, 0, 0], [0, 0, 0], [0, 0, 1]]) / np.sqrt(2)
     truth = ((0.5, np.array([1, 0, 0])), (0.5, np.array([0, 0, 1])))
     for q1 in [0, 1]:
@@ -597,14 +627,14 @@ def test_partial_trace_of_state_vector_as_mixture_mixed_result_qudits():
         assert mixtures_equal(mixture, truth)
 
 
-def test_to_special():
+def test_to_special() -> None:
     u = cirq.testing.random_unitary(4)
     su = cirq.to_special(u)
     assert not cirq.is_special_unitary(u)
     assert cirq.is_special_unitary(su)
 
 
-def test_default_tolerance():
+def test_default_tolerance() -> None:
     a, b = cirq.LineQubit.range(2)
     final_state_vector = (
         cirq.Simulator()
@@ -618,7 +648,7 @@ def test_default_tolerance():
 
 @pytest.mark.parametrize('state_1', [0, 1])
 @pytest.mark.parametrize('state_2', [0, 1])
-def test_factor_state_vector(state_1: int, state_2: int):
+def test_factor_state_vector(state_1: int, state_2: int) -> None:
     # Kron two state vectors and apply a phase. Factoring should produce the expected results.
     n = 12
     for i in range(n):
@@ -636,7 +666,7 @@ def test_factor_state_vector(state_1: int, state_2: int):
 
 
 @pytest.mark.parametrize('num_dimensions', [*range(1, 7)])
-def test_transpose_flattened_array(num_dimensions):
+def test_transpose_flattened_array(num_dimensions) -> None:
     np.random.seed(0)
     for _ in range(10):
         shape = np.random.randint(1, 5, (num_dimensions,)).tolist()
@@ -656,7 +686,7 @@ def test_can_numpy_support_shape(shape: tuple[int, ...], result: bool) -> None:
 
 
 @pytest.mark.parametrize('coeff', [1, 1j, -1, -1j, 1j**0.5, 1j**0.3])
-def test_phase_delta(coeff):
+def test_phase_delta(coeff) -> None:
     u1 = cirq.testing.random_unitary(4)
     u2 = u1 * coeff
     np.testing.assert_almost_equal(linalg.phase_delta(u1, u2), coeff)

cirq/neutral_atoms/convert_to_neutral_atom_gates.py

@@ -12,15 +12,21 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from cirq import ops
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
 from cirq.neutral_atoms import neutral_atom_devices
 
+if TYPE_CHECKING:
+    import cirq
+
 
-def is_native_neutral_atom_op(operation: ops.Operation) -> bool:
+def is_native_neutral_atom_op(operation: cirq.Operation) -> bool:
     """Returns true if the operation is in the default neutral atom gateset."""
     return operation in neutral_atom_devices.neutral_atom_gateset()
 
 
-def is_native_neutral_atom_gate(gate: ops.Gate) -> bool:
+def is_native_neutral_atom_gate(gate: cirq.Gate) -> bool:
     """Returns true if the gate is in the default neutral atom gateset."""
     return gate in neutral_atom_devices.neutral_atom_gateset()

cirq/neutral_atoms/convert_to_neutral_atom_gates_test.py

@@ -12,6 +12,8 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+from __future__ import annotations
+
 import pytest
 
 import cirq
@@ -43,7 +45,7 @@ Q, Q2, Q3 = cirq.LineQubit.range(3)
         (cirq.ZPowGate(exponent=0.5)(Q).controlled_by(Q2, Q3), False),
     ],
 )
-def test_gateset(op: cirq.Operation, expected: bool):
+def test_gateset(op: cirq.Operation, expected: bool) -> None:
     assert cirq.is_native_neutral_atom_op(op) == expected
     if op.gate is not None:
         assert cirq.is_native_neutral_atom_gate(op.gate) == expected

cirq/neutral_atoms/neutral_atom_devices.py

@@ -12,6 +12,8 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+from __future__ import annotations
+
 from cirq import ops
 
 

cirq/ops/__init__.py

@@ -17,6 +17,8 @@ from cirq.ops.arithmetic_operation import ArithmeticGate as ArithmeticGate
 
 from cirq.ops.clifford_gate import (
     CliffordGate as CliffordGate,
+    CXSWAP as CXSWAP,
+    CZSWAP as CZSWAP,
     SingleQubitCliffordGate as SingleQubitCliffordGate,
 )
 

cirq/ops/arithmetic_operation.py

@@ -13,9 +13,11 @@
 # limitations under the License.
 """Helper class for implementing classical arithmetic operations."""
 
+from __future__ import annotations
+
 import abc
 import itertools
-from typing import cast, Iterable, List, Sequence, Tuple, TYPE_CHECKING, Union
+from typing import cast, Iterable, Sequence, TYPE_CHECKING
 
 import numpy as np
 from typing_extensions import Self
@@ -42,21 +44,21 @@ class ArithmeticGate(Gate, metaclass=abc.ABCMeta):
     >>> class Add(cirq.ArithmeticGate):
     ...     def __init__(
     ...         self,
-    ...         target_register: '[int, Sequence[int]]',
-    ...         input_register: 'Union[int, Sequence[int]]',
+    ...         target_register: int | Sequence[int],
+    ...         input_register: int | Sequence[int],
     ...     ):
     ...         self.target_register = target_register
     ...         self.input_register = input_register
     ...
-    ...     def registers(self) -> 'Sequence[Union[int, Sequence[int]]]':
+    ...     def registers(self) -> Sequence[int | Sequence[int]]:
     ...         return self.target_register, self.input_register
     ...
     ...     def with_registers(
-    ...         self, *new_registers: 'Union[int, Sequence[int]]'
+    ...         self, *new_registers: int | Sequence[int]
     ...     ) -> 'Add':
     ...         return Add(*new_registers)
     ...
-    ...     def apply(self, *register_values: int) -> 'Union[int, Iterable[int]]':
+    ...     def apply(self, *register_values: int) -> int | Iterable[int]:
     ...         return sum(register_values)
     >>> cirq.unitary(
     ...     Add(target_register=[2, 2],
@@ -86,7 +88,7 @@ class ArithmeticGate(Gate, metaclass=abc.ABCMeta):
     """
 
     @abc.abstractmethod
-    def registers(self) -> Sequence[Union[int, Sequence[int]]]:
+    def registers(self) -> Sequence[int | Sequence[int]]:
         """The data acted upon by the arithmetic gate.
 
         Each register in the list can either be a classical constant (an `int`),
@@ -103,7 +105,7 @@ class ArithmeticGate(Gate, metaclass=abc.ABCMeta):
         raise NotImplementedError()
 
     @abc.abstractmethod
-    def with_registers(self, *new_registers: Union[int, Sequence[int]]) -> Self:
+    def with_registers(self, *new_registers: int | Sequence[int]) -> Self:
         """Returns the same fate targeting different registers.
 
         Args:
@@ -117,7 +119,7 @@ class ArithmeticGate(Gate, metaclass=abc.ABCMeta):
         raise NotImplementedError()
 
     @abc.abstractmethod
-    def apply(self, *register_values: int) -> Union[int, Iterable[int]]:
+    def apply(self, *register_values: int) -> int | Iterable[int]:
         """Returns the result of the gate operating on classical values.
 
         For example, an addition takes two values (the target and the source),
@@ -162,7 +164,7 @@ class ArithmeticGate(Gate, metaclass=abc.ABCMeta):
         """
         raise NotImplementedError()
 
-    def _qid_shape_(self) -> Tuple[int, ...]:
+    def _qid_shape_(self) -> tuple[int, ...]:
         shape = []
         for r in self.registers():
             if isinstance(r, Sequence):
@@ -170,11 +172,11 @@ class ArithmeticGate(Gate, metaclass=abc.ABCMeta):
                     shape.append(i)
         return tuple(shape)
 
-    def _apply_unitary_(self, args: 'cirq.ApplyUnitaryArgs'):
+    def _apply_unitary_(self, args: cirq.ApplyUnitaryArgs):
         registers = self.registers()
-        input_ranges: List[Sequence[int]] = []
-        shape: List[int] = []
-        overflow_sizes: List[int] = []
+        input_ranges: list[Sequence[int]] = []
+        shape: list[int] = []
+        overflow_sizes: list[int] = []
         for register in registers:
             if isinstance(register, int):
                 input_ranges.append([register])
@@ -196,8 +198,8 @@ class ArithmeticGate(Gate, metaclass=abc.ABCMeta):
             output = self.apply(*input_seq)
 
             # Wrap into list.
-            inputs: List[int] = list(input_seq)
-            outputs: List[int] = [output] if isinstance(output, int) else list(output)
+            inputs: list[int] = list(input_seq)
+            outputs: list[int] = [output] if isinstance(output, int) else list(output)
 
             # Omitted tail values default to the corresponding input value.
             if len(outputs) < len(inputs):
@@ -219,8 +221,8 @@ class ArithmeticGate(Gate, metaclass=abc.ABCMeta):
                     outputs[i] %= overflow_sizes[i]
 
             # Copy amplitude to new location.
-            cast(List[Union[int, slice]], outputs).append(slice(None))
-            cast(List[Union[int, slice]], inputs).append(slice(None))
+            cast(list[int | slice], outputs).append(slice(None))
+            cast(list[int | slice], inputs).append(slice(None))
             dst[tuple(outputs)] = src[tuple(inputs)]
 
         # In case the reshaped arrays were copies instead of views.
@@ -231,9 +233,7 @@ class ArithmeticGate(Gate, metaclass=abc.ABCMeta):
 
 
 def _describe_bad_arithmetic_changed_const(
-    registers: Sequence[Union[int, Sequence[Union['cirq.Qid', int]]]],
-    inputs: List[int],
-    outputs: List[int],
+    registers: Sequence[int | Sequence[cirq.Qid | int]], inputs: list[int], outputs: list[int]
 ) -> str:
     from cirq.circuits import TextDiagramDrawer
 

cirq/ops/arithmetic_operation_test.py

@@ -12,7 +12,9 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from typing import Sequence, Union
+from __future__ import annotations
+
+from typing import Sequence
 
 import numpy as np
 import pytest
@@ -32,11 +34,10 @@ def adder_matrix(target_width: int, source_width: int) -> np.ndarray:
     result = np.zeros((t, s, t, s))
     for k in range(s):
         result[:, k, :, k] = shift_matrix(t, k)
-    result.shape = (t * s, t * s)
-    return result
+    return result.reshape(t * s, t * s)
 
 
-def test_the_tests():
+def test_the_tests() -> None:
     # fmt: off
     np.testing.assert_allclose(
         shift_matrix(4, 1),
@@ -85,12 +86,10 @@ def test_the_tests():
     )
 
 
-def test_arithmetic_gate_apply_unitary():
+def test_arithmetic_gate_apply_unitary() -> None:
     class Add(cirq.ArithmeticGate):
         def __init__(
-            self,
-            target_register: Union[int, Sequence[int]],
-            input_register: Union[int, Sequence[int]],
+            self, target_register: int | Sequence[int], input_register: int | Sequence[int]
         ):
             self.target_register = target_register
             self.input_register = input_register