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

cirq/protocols/kraus_protocol.py

@@ -14,13 +14,16 @@
 
 """Protocol and methods for obtaining Kraus representation of quantum channels."""
 
+from __future__ import annotations
+
 import warnings
 from types import NotImplementedType
-from typing import Any, Sequence, Tuple, TypeVar, Union
+from typing import Any, Sequence, TypeVar
 
 import numpy as np
 from typing_extensions import Protocol
 
+from cirq import protocols, qis
 from cirq._doc import doc_private
 from cirq.protocols.decompose_protocol import _try_decompose_into_operations_and_qubits
 from cirq.protocols.mixture_protocol import has_mixture
@@ -31,7 +34,7 @@ from cirq.protocols.unitary_protocol import unitary
 # Sequence[np.ndarray] to ensure the method has the correct type signature in
 # that case. It is checked for using `is`, so it won't have a false positive
 # if the user provides a different (np.array([]),) value.
-RaiseTypeErrorIfNotProvided: Tuple[np.ndarray] = (np.array([]),)
+RaiseTypeErrorIfNotProvided: tuple[np.ndarray] = (np.array([]),)
 
 
 TDefault = TypeVar('TDefault')
@@ -41,7 +44,7 @@ class SupportsKraus(Protocol):
     """An object that may be describable as a quantum channel."""
 
     @doc_private
-    def _kraus_(self) -> Union[Sequence[np.ndarray], NotImplementedType]:
+    def _kraus_(self) -> Sequence[np.ndarray] | NotImplementedType:
         r"""A list of Kraus matrices describing the quantum channel.
 
         These matrices are the terms in the operator sum representation of a
@@ -92,9 +95,40 @@ class SupportsKraus(Protocol):
         """
 
 
+def _strat_kraus_from_apply_channel(val: Any) -> tuple[np.ndarray, ...] | None:
+    """Attempts to compute a value's Kraus operators via its _apply_channel_ method.
+    This is very expensive (O(16^N)), so only do this as a last resort."""
+    method = getattr(val, '_apply_channel_', None)
+    if method is None:
+        return None
+
+    qid_shape = protocols.qid_shape(val)
+
+    eye = qis.eye_tensor(qid_shape * 2, dtype=np.complex128)
+    buffer = np.empty_like(eye)
+    buffer.fill(float('nan'))
+    superop = protocols.apply_channel(
+        val=val,
+        args=protocols.ApplyChannelArgs(
+            target_tensor=eye,
+            out_buffer=buffer,
+            auxiliary_buffer0=buffer.copy(),
+            auxiliary_buffer1=buffer.copy(),
+            left_axes=list(range(len(qid_shape))),
+            right_axes=list(range(len(qid_shape), len(qid_shape) * 2)),
+        ),
+        default=None,
+    )
+    if superop is None or superop is NotImplemented:
+        return None
+    n = np.prod(qid_shape) ** 2
+    kraus_ops = qis.superoperator_to_kraus(superop.reshape((n, n)))
+    return tuple(kraus_ops)
+
+
 def kraus(
     val: Any, default: Any = RaiseTypeErrorIfNotProvided
-) -> Union[Tuple[np.ndarray, ...], TDefault]:
+) -> tuple[np.ndarray, ...] | TDefault:
     r"""Returns a list of matrices describing the channel for the given value.
 
     These matrices are the terms in the operator sum representation of
@@ -146,9 +180,7 @@ def kraus(
     mixture_getter = getattr(val, '_mixture_', None)
     mixture_result = NotImplemented if mixture_getter is None else mixture_getter()
     if mixture_result is not NotImplemented and mixture_result is not None:
-        return tuple(
-            np.sqrt(p) * (u if isinstance(u, np.ndarray) else unitary(u)) for p, u in mixture_result
-        )
+        return tuple(np.sqrt(p) * unitary(u) for p, u in mixture_result)
 
     unitary_getter = getattr(val, '_unitary_', None)
     unitary_result = NotImplemented if unitary_getter is None else unitary_getter()
@@ -159,6 +191,14 @@ def kraus(
     if channel_result is not NotImplemented:
         return tuple(channel_result)  # pragma: no cover
 
+    # Last-resort fallback: try to derive Kraus from _apply_channel_.
+    # Note: _apply_channel can lead to kraus being called again, so if default
+    # is None, this can trigger an infinite loop.
+    if default is not None:
+        result = _strat_kraus_from_apply_channel(val)
+        if result is not None:
+            return result
+
     if default is not RaiseTypeErrorIfNotProvided:
         return default
 

cirq/protocols/kraus_protocol_test.py

@@ -14,17 +14,20 @@
 
 """Tests for kraus_protocol.py."""
 
-from typing import Iterable, List, Sequence, Tuple
+from __future__ import annotations
+
+from typing import Iterable, Sequence
 
 import numpy as np
 import pytest
 
 import cirq
+from cirq.protocols.apply_channel_protocol import _apply_kraus
 
-LOCAL_DEFAULT: List[np.ndarray] = [np.array([])]
+LOCAL_DEFAULT: list[np.ndarray] = [np.array([])]
 
 
-def test_kraus_no_methods():
+def test_kraus_no_methods() -> None:
     class NoMethod:
         pass
 
@@ -51,7 +54,7 @@ def assert_not_implemented(val):
     assert not cirq.has_kraus(val)
 
 
-def test_kraus_returns_not_implemented():
+def test_kraus_returns_not_implemented() -> None:
     class ReturnsNotImplemented:
         def _kraus_(self):
             return NotImplemented
@@ -59,7 +62,7 @@ def test_kraus_returns_not_implemented():
     assert_not_implemented(ReturnsNotImplemented())
 
 
-def test_mixture_returns_not_implemented():
+def test_mixture_returns_not_implemented() -> None:
     class ReturnsNotImplemented:
         def _mixture_(self):
             return NotImplemented
@@ -67,7 +70,7 @@ def test_mixture_returns_not_implemented():
     assert_not_implemented(ReturnsNotImplemented())
 
 
-def test_unitary_returns_not_implemented():
+def test_unitary_returns_not_implemented() -> None:
     class ReturnsNotImplemented:
         def _unitary_(self):
             return NotImplemented
@@ -80,7 +83,7 @@ def test_unitary_returns_not_implemented():
     assert cirq.kraus(ReturnsNotImplemented(), LOCAL_DEFAULT) is LOCAL_DEFAULT
 
 
-def test_explicit_kraus():
+def test_explicit_kraus() -> None:
     a0 = np.array([[0, 0], [1, 0]])
     a1 = np.array([[1, 0], [0, 0]])
     c = (a0, a1)
@@ -98,11 +101,11 @@ def test_explicit_kraus():
     assert cirq.has_kraus(ReturnsKraus())
 
 
-def test_kraus_fallback_to_mixture():
+def test_kraus_fallback_to_mixture() -> None:
     m = ((0.3, cirq.unitary(cirq.X)), (0.4, cirq.unitary(cirq.Y)), (0.3, cirq.unitary(cirq.Z)))
 
     class ReturnsMixture:
-        def _mixture_(self) -> Iterable[Tuple[float, np.ndarray]]:
+        def _mixture_(self) -> Iterable[tuple[float, np.ndarray]]:
             return m
 
     c = (
@@ -120,7 +123,7 @@ def test_kraus_fallback_to_mixture():
     assert cirq.has_kraus(ReturnsMixture())
 
 
-def test_kraus_fallback_to_unitary():
+def test_kraus_fallback_to_unitary() -> None:
     u = np.array([[1, 0], [1, 0]])
 
     class ReturnsUnitary:
@@ -160,12 +163,83 @@ class HasKrausWhenDecomposed(cirq.testing.SingleQubitGate):
 
 
 @pytest.mark.parametrize('cls', [HasKraus, HasMixture, HasUnitary])
-def test_has_kraus(cls):
+def test_has_kraus(cls) -> None:
     assert cirq.has_kraus(cls())
 
 
 @pytest.mark.parametrize('decomposed_cls', [HasKraus, HasMixture, HasUnitary])
-def test_has_kraus_when_decomposed(decomposed_cls):
+def test_has_kraus_when_decomposed(decomposed_cls) -> None:
     op = HasKrausWhenDecomposed(decomposed_cls).on(cirq.NamedQubit('test'))
     assert cirq.has_kraus(op)
     assert not cirq.has_kraus(op, allow_decompose=False)
+
+
+def test_strat_kraus_from_apply_channel_returns_none():
+    # Remove _kraus_ and _apply_channel_ methods
+    class NoApplyChannelReset(cirq.ResetChannel):
+        def _kraus_(self):
+            return NotImplemented
+
+        def _apply_channel_(self, args):
+            return NotImplemented
+
+    gate_no_apply = NoApplyChannelReset()
+    with pytest.raises(
+        TypeError,
+        match="does have a _kraus_, _mixture_ or _unitary_ method, but it returned NotImplemented",
+    ):
+        cirq.kraus(gate_no_apply)
+
+
+@pytest.mark.parametrize(
+    'channel_cls,params',
+    [
+        (cirq.BitFlipChannel, (0.5,)),
+        (cirq.PhaseFlipChannel, (0.3,)),
+        (cirq.DepolarizingChannel, (0.2,)),
+        (cirq.AmplitudeDampingChannel, (0.4,)),
+        (cirq.PhaseDampingChannel, (0.25,)),
+    ],
+)
+def test_kraus_fallback_to_apply_channel(channel_cls, params) -> None:
+    """Kraus protocol falls back to _apply_channel_ when no _kraus_, _mixture_, or _unitary_."""
+    # Create the expected channel and get its Kraus operators
+    expected_channel = channel_cls(*params)
+    expected_kraus = cirq.kraus(expected_channel)
+
+    class TestChannel:
+        def __init__(self, channel_cls, params):
+            self.channel_cls = channel_cls
+            self.params = params
+            self.expected_kraus = cirq.kraus(channel_cls(*params))
+
+        def _num_qubits_(self):
+            return 1
+
+        def _apply_channel_(self, args: cirq.ApplyChannelArgs):
+            return _apply_kraus(self.expected_kraus, args)
+
+    chan = TestChannel(channel_cls, params)
+    kraus_ops = cirq.kraus(chan)
+
+    # Compare the superoperator matrices for equivalence
+    expected_super = sum(np.kron(k, k.conj()) for k in expected_kraus)
+    actual_super = sum(np.kron(k, k.conj()) for k in kraus_ops)
+    np.testing.assert_allclose(actual_super, expected_super, atol=1e-8)
+
+
+def test_reset_channel_kraus_apply_channel_consistency():
+    Reset = cirq.ResetChannel
+    # Original gate
+    gate = Reset()
+    cirq.testing.assert_has_consistent_apply_channel(gate)
+    cirq.testing.assert_consistent_channel(gate)
+
+    # Remove _kraus_ method
+    class NoKrausReset(Reset):
+        def _kraus_(self):
+            return NotImplemented
+
+    gate_no_kraus = NoKrausReset()
+    # Should still match the original superoperator
+    np.testing.assert_allclose(cirq.kraus(gate), cirq.kraus(gate_no_kraus), atol=1e-8)

cirq/protocols/measurement_key_protocol.py

@@ -11,10 +11,13 @@
 # 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.
+
 """Protocol for object that have measurement keys."""
 
+from __future__ import annotations
+
 from types import NotImplementedType
-from typing import Any, FrozenSet, Mapping, Optional, Tuple, TYPE_CHECKING, Union
+from typing import Any, Mapping, TYPE_CHECKING
 
 from typing_extensions import Protocol
 
@@ -60,7 +63,7 @@ class SupportsMeasurementKey(Protocol):
         """Return if this object is (or contains) a measurement."""
 
     @doc_private
-    def _measurement_key_obj_(self) -> 'cirq.MeasurementKey':
+    def _measurement_key_obj_(self) -> cirq.MeasurementKey:
         """Return the key object that will be used to identify this measurement.
 
         When a measurement occurs, either on hardware, or in a simulation,
@@ -69,9 +72,7 @@ class SupportsMeasurementKey(Protocol):
         """
 
     @doc_private
-    def _measurement_key_objs_(
-        self,
-    ) -> Union[FrozenSet['cirq.MeasurementKey'], NotImplementedType, None]:
+    def _measurement_key_objs_(self) -> frozenset[cirq.MeasurementKey] | NotImplementedType | None:
         """Return the key objects for measurements performed by the receiving object.
 
         When a measurement occurs, either on hardware, or in a simulation,
@@ -89,7 +90,7 @@ class SupportsMeasurementKey(Protocol):
         """
 
     @doc_private
-    def _measurement_key_names_(self) -> Union[FrozenSet[str], NotImplementedType, None]:
+    def _measurement_key_names_(self) -> frozenset[str] | NotImplementedType | None:
         """Return the string keys for measurements performed by the receiving object.
 
         When a measurement occurs, either on hardware, or in a simulation,
@@ -175,7 +176,7 @@ def measurement_key_name(val: Any, default: Any = RaiseTypeErrorIfNotProvided):
 
 def _measurement_key_objs_from_magic_methods(
     val: Any,
-) -> Union[FrozenSet['cirq.MeasurementKey'], NotImplementedType, None]:
+) -> frozenset[cirq.MeasurementKey] | NotImplementedType | None:
     """Uses the measurement key related magic methods to get the `MeasurementKey`s for this
     object."""
 
@@ -193,7 +194,7 @@ def _measurement_key_objs_from_magic_methods(
 
 def _measurement_key_names_from_magic_methods(
     val: Any,
-) -> Union[FrozenSet[str], NotImplementedType, None]:
+) -> frozenset[str] | NotImplementedType | None:
     """Uses the measurement key related magic methods to get the key strings for this object."""
 
     getter = getattr(val, '_measurement_key_names_', None)
@@ -209,7 +210,7 @@ def _measurement_key_names_from_magic_methods(
     return result
 
 
-def measurement_key_objs(val: Any) -> FrozenSet['cirq.MeasurementKey']:
+def measurement_key_objs(val: Any) -> frozenset[cirq.MeasurementKey]:
     """Gets the measurement key objects of measurements within the given value.
 
     Args:
@@ -228,7 +229,7 @@ def measurement_key_objs(val: Any) -> FrozenSet['cirq.MeasurementKey']:
     return frozenset()
 
 
-def measurement_key_names(val: Any) -> FrozenSet[str]:
+def measurement_key_names(val: Any) -> frozenset[str]:
     """Gets the measurement key strings of measurements within the given value.
 
     Args:
@@ -253,7 +254,7 @@ def measurement_key_names(val: Any) -> FrozenSet[str]:
     return frozenset()
 
 
-def _is_measurement_from_magic_method(val: Any) -> Optional[bool]:
+def _is_measurement_from_magic_method(val: Any) -> bool | None:
     """Uses `is_measurement` magic method to determine if this object is a measurement."""
     getter = getattr(val, '_is_measurement_', None)
     return NotImplemented if getter is None else getter()
@@ -289,7 +290,7 @@ def with_measurement_key_mapping(val: Any, key_map: Mapping[str, str]):
     return NotImplemented if getter is None else getter(key_map)
 
 
-def with_key_path(val: Any, path: Tuple[str, ...]):
+def with_key_path(val: Any, path: tuple[str, ...]):
     """Adds the path to the target's measurement keys.
 
     The path usually refers to an identifier or a list of identifiers from a subcircuit that
@@ -300,7 +301,7 @@ def with_key_path(val: Any, path: Tuple[str, ...]):
     return NotImplemented if getter is None else getter(path)
 
 
-def with_key_path_prefix(val: Any, prefix: Tuple[str, ...]):
+def with_key_path_prefix(val: Any, prefix: tuple[str, ...]):
     """Prefixes the path to the target's measurement keys.
 
     The path usually refers to an identifier or a list of identifiers from a subcircuit that
@@ -316,9 +317,7 @@ def with_key_path_prefix(val: Any, prefix: Tuple[str, ...]):
 
 
 def with_rescoped_keys(
-    val: Any,
-    path: Tuple[str, ...],
-    bindable_keys: Optional[FrozenSet['cirq.MeasurementKey']] = None,
+    val: Any, path: tuple[str, ...], bindable_keys: frozenset[cirq.MeasurementKey] | None = None
 ):
     """Rescopes any measurement and control keys to the provided path, given the existing keys.
 

cirq/protocols/measurement_key_protocol_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
@@ -28,7 +30,7 @@ class ReturnsObj:
 
 
 @pytest.mark.parametrize('gate', [ReturnsStr(), ReturnsObj()])
-def test_measurement_key_name(gate):
+def test_measurement_key_name(gate) -> None:
     assert isinstance(cirq.measurement_key_name(gate), str)
     assert cirq.measurement_key_name(gate) == 'door locker'
     assert cirq.measurement_key_obj(gate) == cirq.MeasurementKey(name='door locker')
@@ -39,7 +41,7 @@ def test_measurement_key_name(gate):
 
 
 @pytest.mark.parametrize('gate', [ReturnsStr(), ReturnsObj()])
-def test_measurement_key_obj(gate):
+def test_measurement_key_obj(gate) -> None:
     assert isinstance(cirq.measurement_key_obj(gate), cirq.MeasurementKey)
     assert cirq.measurement_key_obj(gate) == cirq.MeasurementKey(name='door locker')
     assert cirq.measurement_key_obj(gate) == 'door locker'
@@ -50,7 +52,7 @@ def test_measurement_key_obj(gate):
 
 
 @pytest.mark.parametrize('key_method', [cirq.measurement_key_name, cirq.measurement_key_obj])
-def test_measurement_key_no_method(key_method):
+def test_measurement_key_no_method(key_method) -> None:
     class NoMethod:
         pass
 
@@ -73,7 +75,7 @@ def test_measurement_key_no_method(key_method):
 
 
 @pytest.mark.parametrize('key_method', [cirq.measurement_key_name, cirq.measurement_key_obj])
-def test_measurement_key_not_implemented_default_behavior(key_method):
+def test_measurement_key_not_implemented_default_behavior(key_method) -> None:
     class ReturnsNotImplemented:
         def _measurement_key_name_(self):
             return NotImplemented
@@ -89,7 +91,7 @@ def test_measurement_key_not_implemented_default_behavior(key_method):
     assert key_method(ReturnsNotImplemented(), 'a') == 'a'
 
 
-def test_is_measurement():
+def test_is_measurement() -> None:
     q = cirq.NamedQubit('q')
     assert cirq.is_measurement(cirq.measure(q))
     assert cirq.is_measurement(cirq.MeasurementGate(num_qubits=1, key='b'))
@@ -99,7 +101,7 @@ def test_is_measurement():
     assert not cirq.is_measurement(cirq.bit_flip(1))
 
     class NotImplementedOperation(cirq.Operation):
-        def with_qubits(self, *new_qubits) -> 'NotImplementedOperation':
+        def with_qubits(self, *new_qubits) -> NotImplementedOperation:
             raise NotImplementedError()
 
         @property
@@ -109,7 +111,7 @@ def test_is_measurement():
     assert not cirq.is_measurement(NotImplementedOperation())
 
 
-def test_measurement_without_key():
+def test_measurement_without_key() -> None:
     class MeasurementWithoutKey:
         def _is_measurement_(self):
             return True
@@ -120,7 +122,7 @@ def test_measurement_without_key():
     assert cirq.is_measurement(MeasurementWithoutKey())
 
 
-def test_non_measurement_with_key():
+def test_non_measurement_with_key() -> None:
     class NonMeasurementGate(cirq.Gate):
         def _is_measurement_(self):
             return False  # pragma: no cover
@@ -155,7 +157,7 @@ def test_non_measurement_with_key():
     ('key_method', 'keys'),
     [(cirq.measurement_key_names, {'a', 'b'}), (cirq.measurement_key_objs, {'c', 'd'})],
 )
-def test_measurement_keys(key_method, keys):
+def test_measurement_keys(key_method, keys) -> None:
     class MeasurementKeysGate(cirq.Gate):
         def _measurement_key_names_(self):
             return frozenset(['a', 'b'])
@@ -180,7 +182,7 @@ def test_measurement_keys(key_method, keys):
     assert key_method(MeasurementKeysGate().on(a)) == keys
 
 
-def test_measurement_key_mapping():
+def test_measurement_key_mapping() -> None:
     class MultiKeyGate:
         def __init__(self, keys):
             self._keys = frozenset(keys)
@@ -217,7 +219,7 @@ def test_measurement_key_mapping():
     assert cirq.measurement_key_names(mkg_cdx) == {'c', 'd'}
 
 
-def test_measurement_key_path():
+def test_measurement_key_path() -> None:
     class MultiKeyGate:
         def __init__(self, keys):
             self._keys = frozenset(cirq.MeasurementKey.parse_serialized(key) for key in keys)

cirq/protocols/mixture_protocol.py

@@ -14,8 +14,10 @@
 
 """Protocol for objects that are mixtures (probabilistic combinations)."""
 
+from __future__ import annotations
+
 from types import NotImplementedType
-from typing import Any, Sequence, Tuple, Union
+from typing import Any, Sequence
 
 import numpy as np
 from typing_extensions import Protocol
@@ -27,14 +29,14 @@ from cirq.protocols.unitary_protocol import unitary
 
 # This is a special indicator value used by the inverse method to determine
 # whether or not the caller provided a 'default' argument.
-RaiseTypeErrorIfNotProvided: Sequence[Tuple[float, Any]] = ((0.0, []),)
+RaiseTypeErrorIfNotProvided: Sequence[tuple[float, Any]] = ((0.0, []),)
 
 
 class SupportsMixture(Protocol):
     """An object that decomposes into a probability distribution of unitaries."""
 
     @doc_private
-    def _mixture_(self) -> Union[Sequence[Tuple[float, Any]], NotImplementedType]:
+    def _mixture_(self) -> Sequence[tuple[float, Any]] | NotImplementedType:
         """Decompose into a probability distribution of unitaries.
 
         This method is used by the global `cirq.mixture` method.
@@ -65,7 +67,7 @@ class SupportsMixture(Protocol):
 
 def mixture(
     val: Any, default: Any = RaiseTypeErrorIfNotProvided
-) -> Sequence[Tuple[float, np.ndarray]]:
+) -> Sequence[tuple[float, np.ndarray]]:
     """Return a sequence of tuples representing a probabilistic unitary.
 
     A mixture is described by an iterable of tuples of the form
@@ -92,7 +94,7 @@ def mixture(
     mixture_getter = getattr(val, '_mixture_', None)
     result = NotImplemented if mixture_getter is None else mixture_getter()
     if result is not NotImplemented and result is not None:
-        return tuple((p, u if isinstance(u, np.ndarray) else unitary(u)) for p, u in result)
+        return tuple((p, unitary(u)) for p, u in result)
 
     unitary_getter = getattr(val, '_unitary_', None)
     result = NotImplemented if unitary_getter is None else unitary_getter()

cirq/protocols/mixture_protocol_test.py

@@ -12,13 +12,15 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+from __future__ import annotations
+
 import numpy as np
 import pytest
 
 import cirq
 
-a = np.array([1])
-b = np.array([1j])
+a = np.array([[1]])
+b = np.array([[1j]])
 
 
 class NoMethod:

cirq/protocols/mul_protocol.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
 
 from cirq.protocols.resolve_parameters import is_parameterized
@@ -69,6 +71,3 @@ def mul(lhs: Any, rhs: Any, default: Any = RaiseTypeErrorIfNotProvided) -> Any:
     if default is not RaiseTypeErrorIfNotProvided:
         return default
     raise TypeError(f"unsupported operand type(s) for *: '{type(lhs)}' and '{type(rhs)}'")
-
-
-# pylint: enable=function-redefined, redefined-builtin

cirq/protocols/mul_protocol_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 sympy
 

cirq/protocols/pauli_expansion_protocol.py

@@ -11,9 +11,12 @@
 # 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.
+
 """Protocol for obtaining expansion of linear operators in Pauli basis."""
 
-from typing import Any, TypeVar, Union
+from __future__ import annotations
+
+from typing import Any, TypeVar
 
 from typing_extensions import Protocol
 
@@ -45,9 +48,9 @@ class SupportsPauliExpansion(Protocol):
 def pauli_expansion(
     val: Any,
     *,
-    default: Union[value.LinearDict[str], TDefault] = RaiseTypeErrorIfNotProvided,
+    default: value.LinearDict[str] | TDefault = RaiseTypeErrorIfNotProvided,
     atol: float = 1e-9,
-) -> Union[value.LinearDict[str], TDefault]:
+) -> value.LinearDict[str] | TDefault:
     """Returns coefficients of the expansion of val in the Pauli basis.
 
     Args:

cirq/protocols/pauli_expansion_protocol_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
 
@@ -52,9 +54,9 @@ class HasQuditUnitary:
 
 
 @pytest.mark.parametrize(
-    'val', (NoMethod(), ReturnsNotImplemented(), HasQuditUnitary(), 123, np.eye(2), object(), cirq)
+    'val', (NoMethod(), ReturnsNotImplemented(), HasQuditUnitary(), 123, object(), cirq)
 )
-def test_raises_no_pauli_expansion(val):
+def test_raises_no_pauli_expansion(val) -> None:
     assert cirq.pauli_expansion(val, default=None) is None
     with pytest.raises(TypeError, match='No Pauli expansion'):
         cirq.pauli_expansion(val)
@@ -78,7 +80,7 @@ def test_raises_no_pauli_expansion(val):
         ),
     ),
 )
-def test_pauli_expansion(val, expected_expansion):
+def test_pauli_expansion(val, expected_expansion) -> None:
     actual_expansion = cirq.pauli_expansion(val)
     assert cirq.approx_eq(actual_expansion, expected_expansion, atol=1e-12)
     assert set(actual_expansion.keys()) == set(expected_expansion.keys())

cirq/protocols/phase_protocol.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, TypeVar
 
 from typing_extensions import Protocol

cirq/protocols/phase_protocol_test.py

@@ -12,12 +12,14 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+from __future__ import annotations
+
 import pytest
 
 import cirq
 
 
-def test_phase_by():
+def test_phase_by() -> None:
     class NoMethod:
         pass