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

cirq/protocols/apply_mixture_protocol.py

@@ -14,13 +14,14 @@
 
 """A protocol for implementing high performance mixture evolutions."""
 
+from __future__ import annotations
+
 from types import NotImplementedType
-from typing import Any, cast, Iterable, Optional, Tuple, TypeVar, Union
+from typing import Any, cast, Iterable, TypeVar
 
 import numpy as np
 from typing_extensions import Protocol
 
-from cirq import linalg
 from cirq._doc import doc_private
 from cirq.protocols import qid_shape_protocol
 from cirq.protocols.apply_unitary_protocol import apply_unitary, ApplyUnitaryArgs
@@ -73,7 +74,7 @@ class ApplyMixtureArgs:
         auxiliary_buffer0: np.ndarray,
         auxiliary_buffer1: np.ndarray,
         left_axes: Iterable[int],
-        right_axes: Optional[Iterable[int]] = None,
+        right_axes: Iterable[int] | None = None,
     ):
         """Args for apply mixture.
 
@@ -112,9 +113,7 @@ class SupportsApplyMixture(Protocol):
     """An object that can efficiently implement a mixture."""
 
     @doc_private
-    def _apply_mixture_(
-        self, args: ApplyMixtureArgs
-    ) -> Union[np.ndarray, None, NotImplementedType]:
+    def _apply_mixture_(self, args: ApplyMixtureArgs) -> np.ndarray | None | NotImplementedType:
         """Efficiently applies a mixture.
 
         This method is given both the target tensor and workspace of the same
@@ -157,8 +156,8 @@ def apply_mixture(
     val: Any,
     args: ApplyMixtureArgs,
     *,
-    default: Union[np.ndarray, TDefault] = RaiseTypeErrorIfNotProvided,
-) -> Union[np.ndarray, TDefault]:
+    default: np.ndarray | TDefault = RaiseTypeErrorIfNotProvided,
+) -> np.ndarray | TDefault:
     """High performance evolution under a mixture of unitaries evolution.
 
     Follows the steps below to attempt to apply a mixture:
@@ -273,7 +272,7 @@ def apply_mixture(
     )
 
 
-def _validate_input(val: Any, args: 'ApplyMixtureArgs') -> Tuple[Any, 'ApplyMixtureArgs', bool]:
+def _validate_input(val: Any, args: ApplyMixtureArgs) -> tuple[Any, ApplyMixtureArgs, bool]:
     """Validate args input and determine if we are operating on a
     density matrix or a state vector.
     """
@@ -303,8 +302,8 @@ def _validate_input(val: Any, args: 'ApplyMixtureArgs') -> Tuple[Any, 'ApplyMixt
 
 
 def _apply_unitary_strat(
-    val: Any, args: 'ApplyMixtureArgs', is_density_matrix: bool
-) -> Optional[np.ndarray]:
+    val: Any, args: ApplyMixtureArgs, is_density_matrix: bool
+) -> np.ndarray | None:
     """Attempt to use `apply_unitary` and return the result.
 
     If `val` does not support `apply_unitary` returns None.
@@ -325,42 +324,16 @@ def _apply_unitary_strat(
     right_args = ApplyUnitaryArgs(
         target_tensor=np.conjugate(left_result),
         available_buffer=args.auxiliary_buffer0,
-        axes=cast(Tuple[int], args.right_axes),
+        axes=cast(tuple[int], args.right_axes),
     )
     right_result = apply_unitary(val, right_args)
     np.conjugate(right_result, out=right_result)
     return right_result
 
 
-def _apply_unitary_from_matrix_strat(
-    val: np.ndarray, args: 'ApplyMixtureArgs', is_density_matrix: bool
-) -> Optional[np.ndarray]:
-    """Used to enact mixture tuples that are given as (probability, np.ndarray)
-
-    If `val` does not support `apply_unitary` returns None.
-    """
-    qid_shape = tuple(args.target_tensor.shape[i] for i in args.left_axes)
-    matrix_tensor = np.reshape(val.astype(args.target_tensor.dtype), qid_shape * 2)
-    linalg.targeted_left_multiply(
-        matrix_tensor, args.target_tensor, args.left_axes, out=args.auxiliary_buffer0
-    )
-
-    if not is_density_matrix:
-        return args.auxiliary_buffer0
-    # No need to transpose as we are acting on the tensor
-    # representation of matrix, so transpose is done for us.
-    linalg.targeted_left_multiply(
-        np.conjugate(matrix_tensor),
-        args.auxiliary_buffer0,
-        cast(Tuple[int], args.right_axes),
-        out=args.target_tensor,
-    )
-    return args.target_tensor
-
-
 def _apply_mixture_from_mixture_strat(
-    val: Any, args: 'ApplyMixtureArgs', is_density_matrix: bool
-) -> Optional[np.ndarray]:
+    val: Any, args: ApplyMixtureArgs, is_density_matrix: bool
+) -> np.ndarray | None:
     """Attempt to use unitary matrices in _mixture_ and return the result."""
     method = getattr(val, '_mixture_', None)
     if method is None:
@@ -373,8 +346,6 @@ def _apply_mixture_from_mixture_strat(
     for prob, op in prob_mix:
         np.copyto(dst=args.target_tensor, src=args.auxiliary_buffer1)
         right_result = _apply_unitary_strat(op, args, is_density_matrix)
-        if right_result is None:
-            right_result = _apply_unitary_from_matrix_strat(op, args, is_density_matrix)
 
         args.out_buffer += prob * right_result
 

cirq/protocols/apply_mixture_protocol_test.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 Any, cast, Iterable, Optional, Tuple
+from __future__ import annotations
+
+from typing import Any, cast, Iterable
 
 import numpy as np
 import pytest
@@ -31,9 +33,9 @@ def assert_apply_mixture_returns(
     val: Any,
     rho: np.ndarray,
     left_axes: Iterable[int],
-    right_axes: Optional[Iterable[int]],
+    right_axes: Iterable[int] | None,
     assert_result_is_out_buf: bool = False,
-    expected_result: Optional[np.ndarray] = None,
+    expected_result: np.ndarray | None = None,
 ):
     out_buf, buf0, buf1 = make_buffers(rho.shape, rho.dtype)
     result = cirq.apply_mixture(
@@ -94,8 +96,8 @@ def test_apply_mixture_simple():
         def _apply_mixture_(self, args: cirq.ApplyMixtureArgs):
             zero_left = cirq.slice_for_qubits_equal_to(args.left_axes, 0)
             one_left = cirq.slice_for_qubits_equal_to(args.left_axes, 1)
-            zero_right = cirq.slice_for_qubits_equal_to(cast(Tuple[int], args.right_axes), 0)
-            one_right = cirq.slice_for_qubits_equal_to(cast(Tuple[int], args.right_axes), 1)
+            zero_right = cirq.slice_for_qubits_equal_to(cast(tuple[int], args.right_axes), 0)
+            one_right = cirq.slice_for_qubits_equal_to(cast(tuple[int], args.right_axes), 1)
             args.out_buffer[:] = 0
             np.copyto(dst=args.auxiliary_buffer0, src=args.target_tensor)
             for kraus_op in [np.sqrt(0.5) * np.eye(2, dtype=np.complex128), np.sqrt(0.5) * x]:

cirq/protocols/apply_unitary_protocol.py

@@ -14,9 +14,11 @@
 
 """A protocol for implementing high performance unitary left-multiplies."""
 
+from __future__ import annotations
+
 import warnings
-from types import NotImplementedType
-from typing import Any, cast, Iterable, Optional, Sequence, Tuple, TYPE_CHECKING, TypeVar, Union
+from types import EllipsisType, NotImplementedType
+from typing import Any, cast, Iterable, Sequence, TYPE_CHECKING, TypeVar
 
 import numpy as np
 from typing_extensions import Protocol
@@ -73,7 +75,7 @@ class ApplyUnitaryArgs:
         target_tensor: np.ndarray,
         available_buffer: np.ndarray,
         axes: Iterable[int],
-        subspaces: Optional[Sequence[Tuple[int, ...]]] = None,
+        subspaces: Sequence[tuple[int, ...]] | None = None,
     ):
         """Inits ApplyUnitaryArgs.
 
@@ -111,8 +113,8 @@ class ApplyUnitaryArgs:
 
     @staticmethod
     def default(
-        num_qubits: Optional[int] = None, *, qid_shape: Optional[Tuple[int, ...]] = None
-    ) -> 'ApplyUnitaryArgs':
+        num_qubits: int | None = None, *, qid_shape: tuple[int, ...] | None = None
+    ) -> ApplyUnitaryArgs:
         """A default instance starting in state |0⟩.
 
         Specify exactly one argument.
@@ -130,15 +132,15 @@ class ApplyUnitaryArgs:
             raise TypeError('Specify exactly one of num_qubits or qid_shape.')
         if num_qubits is not None:
             qid_shape = (2,) * num_qubits
-        qid_shape = cast(Tuple[int, ...], qid_shape)  # Satisfy mypy
+        qid_shape = cast(tuple[int, ...], qid_shape)  # Satisfy mypy
         num_qubits = len(qid_shape)
         state = qis.one_hot(index=(0,) * num_qubits, shape=qid_shape, dtype=np.complex128)
         return ApplyUnitaryArgs(state, np.empty_like(state), range(num_qubits))
 
     @classmethod
     def for_unitary(
-        cls, num_qubits: Optional[int] = None, *, qid_shape: Optional[Tuple[int, ...]] = None
-    ) -> 'ApplyUnitaryArgs':
+        cls, num_qubits: int | None = None, *, qid_shape: tuple[int, ...] | None = None
+    ) -> ApplyUnitaryArgs:
         """A default instance corresponding to an identity matrix.
 
         Specify exactly one argument.
@@ -157,12 +159,12 @@ class ApplyUnitaryArgs:
             raise TypeError('Specify exactly one of num_qubits or qid_shape.')
         if num_qubits is not None:
             qid_shape = (2,) * num_qubits
-        qid_shape = cast(Tuple[int, ...], qid_shape)  # Satisfy mypy
+        qid_shape = cast(tuple[int, ...], qid_shape)  # Satisfy mypy
         num_qubits = len(qid_shape)
         state = qis.eye_tensor(qid_shape, dtype=np.complex128)
         return ApplyUnitaryArgs(state, np.empty_like(state), range(num_qubits))
 
-    def with_axes_transposed_to_start(self) -> 'ApplyUnitaryArgs':
+    def with_axes_transposed_to_start(self) -> ApplyUnitaryArgs:
         """Returns a transposed view of the same arguments.
 
         Returns:
@@ -180,8 +182,8 @@ class ApplyUnitaryArgs:
         return ApplyUnitaryArgs(target_tensor, available_buffer, range(len(self.axes)))
 
     def _for_operation_with_qid_shape(
-        self, indices: Iterable[int], slices: Tuple[Union[int, slice], ...]
-    ) -> 'ApplyUnitaryArgs':
+        self, indices: Iterable[int], slices: tuple[int | slice, ...]
+    ) -> ApplyUnitaryArgs:
         """Creates a sliced and transposed view of `self` appropriate for an
         operation with shape `qid_shape` on qubits with the given indices.
 
@@ -213,7 +215,7 @@ class ApplyUnitaryArgs:
 
     def subspace_index(
         self, little_endian_bits_int: int = 0, *, big_endian_bits_int: int = 0
-    ) -> Tuple[Union[slice, int, 'ellipsis'], ...]:
+    ) -> tuple[slice | int | EllipsisType, ...]:
         """An index for the subspace where the target axes equal a value.
 
         Args:
@@ -258,9 +260,7 @@ class SupportsConsistentApplyUnitary(Protocol):
     """An object that can be efficiently left-multiplied into tensors."""
 
     @doc_private
-    def _apply_unitary_(
-        self, args: ApplyUnitaryArgs
-    ) -> Union[np.ndarray, None, NotImplementedType]:
+    def _apply_unitary_(self, args: ApplyUnitaryArgs) -> np.ndarray | None | NotImplementedType:
         """Left-multiplies a unitary effect onto a tensor with good performance.
 
         This method is given both the target tensor and workspace of the same
@@ -306,10 +306,10 @@ class SupportsConsistentApplyUnitary(Protocol):
 def apply_unitary(
     unitary_value: Any,
     args: ApplyUnitaryArgs,
-    default: Union[np.ndarray, TDefault] = RaiseTypeErrorIfNotProvided,
+    default: np.ndarray | TDefault = RaiseTypeErrorIfNotProvided,
     *,
     allow_decompose: bool = True,
-) -> Union[np.ndarray, TDefault]:
+) -> np.ndarray | TDefault:
     """High performance left-multiplication of a unitary effect onto a tensor.
 
     Applies the unitary effect of `unitary_value` to the tensor specified in
@@ -420,7 +420,7 @@ def apply_unitary(
 
 def _strat_apply_unitary_from_apply_unitary(
     unitary_value: Any, args: ApplyUnitaryArgs
-) -> Optional[np.ndarray]:
+) -> np.ndarray | None:
     # Check for magic method.
     func = getattr(unitary_value, '_apply_unitary_', None)
     if func is None:
@@ -468,21 +468,26 @@ def _apply_unitary_from_matrix(matrix: np.ndarray, unitary_value: Any, args: App
 
 def _strat_apply_unitary_from_unitary(
     unitary_value: Any, args: ApplyUnitaryArgs
-) -> Optional[np.ndarray]:
-    # Check for magic method.
-    method = getattr(unitary_value, '_unitary_', None)
-    if method is None:
-        return NotImplemented
+) -> np.ndarray | None:
+    if isinstance(unitary_value, np.ndarray):
+        matrix = unitary_value
+        if not linalg.is_unitary(matrix):
+            return None
+    else:
+        # Check for magic method.
+        method = getattr(unitary_value, '_unitary_', None)
+        if method is None:
+            return NotImplemented
 
-    # Attempt to get the unitary matrix.
-    matrix = method()
-    if matrix is NotImplemented or matrix is None:
-        return matrix
+        # Attempt to get the unitary matrix.
+        matrix = method()
+        if matrix is NotImplemented or matrix is None:
+            return matrix
 
     return _apply_unitary_from_matrix(matrix, unitary_value, args)
 
 
-def _strat_apply_unitary_from_decompose(val: Any, args: ApplyUnitaryArgs) -> Optional[np.ndarray]:
+def _strat_apply_unitary_from_decompose(val: Any, args: ApplyUnitaryArgs) -> np.ndarray | None:
     operations, qubits, _ = _try_decompose_into_operations_and_qubits(val)
     if operations is None:
         return NotImplemented
@@ -506,10 +511,10 @@ def _strat_apply_unitary_from_decompose(val: Any, args: ApplyUnitaryArgs) -> Opt
 
 def apply_unitaries(
     unitary_values: Iterable[Any],
-    qubits: Sequence['cirq.Qid'],
-    args: Optional[ApplyUnitaryArgs] = None,
+    qubits: Sequence[cirq.Qid],
+    args: ApplyUnitaryArgs | None = None,
     default: Any = RaiseTypeErrorIfNotProvided,
-) -> Optional[np.ndarray]:
+) -> np.ndarray | None:
     """Apply a series of unitaries onto a state tensor.
 
     Uses `cirq.apply_unitary` on each of the unitary values, to apply them to
@@ -589,7 +594,7 @@ def apply_unitaries(
 
 
 def _incorporate_result_into_target(
-    args: 'ApplyUnitaryArgs', sub_args: 'ApplyUnitaryArgs', sub_result: np.ndarray
+    args: ApplyUnitaryArgs, sub_args: ApplyUnitaryArgs, sub_result: np.ndarray
 ):
     """Takes the result of calling `_apply_unitary_` on `sub_args` and
     copies it back into `args.target_tensor` or `args.available_buffer` as
@@ -645,7 +650,7 @@ def _incorporate_result_into_target(
     return args.target_tensor
 
 
-def _to_slice(subspace_def: Tuple[int, ...]):
+def _to_slice(subspace_def: tuple[int, ...]):
     if len(subspace_def) < 1:
         raise ValueError(f'Subspace {subspace_def} has zero dimensions.')
 

cirq/protocols/apply_unitary_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
 
@@ -54,12 +56,13 @@ def test_apply_unitary_presence_absence():
             args.target_tensor[one] *= -1
             return args.target_tensor
 
-    fails = [NoUnitaryEffect(), HasApplyReturnsNotImplemented()]
+    fails = [NoUnitaryEffect(), HasApplyReturnsNotImplemented(), m * 2]
     passes = [
         HasUnitary(),
         HasApplyReturnsNotImplementedButHasUnitary(),
         HasApplyOutputInBuffer(),
         HasApplyMutateInline(),
+        m,
     ]
 
     def make_input():

cirq/protocols/approximate_equality_protocol.py

@@ -12,6 +12,8 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+from __future__ import annotations
+
 import numbers
 from decimal import Decimal
 from fractions import Fraction

cirq/protocols/approximate_equality_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
+
 from decimal import Decimal
 from fractions import Fraction
 from numbers import Number

cirq/protocols/circuit_diagram_info_protocol.py

@@ -12,20 +12,11 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+from __future__ import annotations
+
 import re
 from fractions import Fraction
-from typing import (
-    Any,
-    Dict,
-    Iterable,
-    List,
-    Optional,
-    overload,
-    Sequence,
-    TYPE_CHECKING,
-    TypeVar,
-    Union,
-)
+from typing import Any, Iterable, overload, Sequence, TYPE_CHECKING, TypeVar, Union
 
 import numpy as np
 import sympy
@@ -50,7 +41,7 @@ class CircuitDiagramInfo:
         wire_symbols: Iterable[str],
         exponent: Any = 1,
         connected: bool = True,
-        exponent_qubit_index: Optional[int] = None,
+        exponent_qubit_index: int | None = None,
         auto_exponent_parens: bool = True,
     ) -> None:
         """Inits CircuitDiagramInfo.
@@ -104,8 +95,8 @@ class CircuitDiagramInfo:
         )
 
     def _wire_symbols_including_formatted_exponent(
-        self, args: 'cirq.CircuitDiagramInfoArgs', *, preferred_exponent_index: Optional[int] = None
-    ) -> List[str]:
+        self, args: cirq.CircuitDiagramInfoArgs, *, preferred_exponent_index: int | None = None
+    ) -> list[str]:
         result = list(self.wire_symbols)
         exponent = self._formatted_exponent(args)
         if exponent is not None:
@@ -122,7 +113,7 @@ class CircuitDiagramInfo:
                 result[k] += f"^{exponent}"
         return result
 
-    def _formatted_exponent(self, args: 'cirq.CircuitDiagramInfoArgs') -> Optional[str]:
+    def _formatted_exponent(self, args: cirq.CircuitDiagramInfoArgs) -> str | None:
         if protocols.is_parameterized(self.exponent):
             name = str(self.exponent)
             return f'({name})' if _is_exposed_formula(name) else name
@@ -190,21 +181,24 @@ class CircuitDiagramInfoArgs:
         precision: The number of digits after the decimal to show for numbers in
             the text diagram. None means use full precision.
         label_map: The map from label entities to diagram positions.
-        include_tags: Whether to print tags from TaggedOperations.
+        include_tags: If ``True`` all tags from ``TaggedOperations`` will be
+            printed.  If ``False`` no tags will be printed.  Alternatively a
+            collection of tag classes can be provided.  In this case only tags
+            whose type is contained in the collection will be shown.
         transpose: Whether the circuit is to be drawn with time from left to
             right (transpose is False), or from top to bottom.
     """
 
-    UNINFORMED_DEFAULT: 'CircuitDiagramInfoArgs'
+    UNINFORMED_DEFAULT: CircuitDiagramInfoArgs
 
     def __init__(
         self,
-        known_qubits: Optional[Iterable['cirq.Qid']],
-        known_qubit_count: Optional[int],
+        known_qubits: Iterable[cirq.Qid] | None,
+        known_qubit_count: int | None,
         use_unicode_characters: bool,
-        precision: Optional[int],
-        label_map: Optional[Dict['cirq.LabelEntity', int]],
-        include_tags: bool = True,
+        precision: int | None,
+        label_map: dict[LabelEntity, int] | None,
+        include_tags: bool | Iterable[type] = True,
         transpose: bool = False,
     ) -> None:
         self.known_qubits = None if known_qubits is None else tuple(known_qubits)
@@ -212,7 +206,11 @@ class CircuitDiagramInfoArgs:
         self.use_unicode_characters = use_unicode_characters
         self.precision = precision
         self.label_map = label_map
-        self.include_tags = include_tags
+        self.include_tags: bool | frozenset[type]
+        if isinstance(include_tags, bool):
+            self.include_tags = include_tags
+        else:
+            self.include_tags = frozenset(include_tags)
         self.transpose = transpose
 
     def _value_equality_values_(self) -> Any:
@@ -226,7 +224,11 @@ class CircuitDiagramInfoArgs:
                 if self.label_map is None
                 else tuple(sorted(self.label_map.items(), key=lambda e: e[0]))
             ),
-            self.include_tags,
+            (
+                self.include_tags
+                if isinstance(self.include_tags, bool)
+                else tuple(sorted(self.include_tags, key=lambda c: c.__name__))
+            ),
             self.transpose,
         )
 
@@ -238,11 +240,28 @@ class CircuitDiagramInfoArgs:
             f'use_unicode_characters={self.use_unicode_characters!r}, '
             f'precision={self.precision!r}, '
             f'label_map={self.label_map!r}, '
-            f'include_tags={self.include_tags!r}, '
+            f'include_tags={self._include_tags_repr()}, '
             f'transpose={self.transpose!r})'
         )
 
-    def format_real(self, val: Union[sympy.Basic, int, float]) -> str:
+    def _include_tags_repr(self) -> str:
+        if isinstance(self.include_tags, bool):
+            return repr(self.include_tags)
+        items = []
+        for cls in self.include_tags:
+            if cls.__module__ == 'builtins':
+                items.append(cls.__qualname__)
+            else:
+                items.append(f"{cls.__module__}.{cls.__qualname__}")
+        joined = ', '.join(items)
+        return f'frozenset({{{joined}}})'
+
+    def tags_to_include(self, tags: Iterable[Any]) -> list[Any]:
+        if isinstance(self.include_tags, bool):
+            return list(tags) if self.include_tags else []
+        return [t for t in tags if any(isinstance(t, cls) for cls in self.include_tags)]
+
+    def format_real(self, val: sympy.Basic | int | float) -> str:
         if isinstance(val, sympy.Basic):
             return str(val)
         if val == int(val):
@@ -251,7 +270,7 @@ class CircuitDiagramInfoArgs:
             return str(val)
         return f'{float(val):.{self.precision}}'
 
-    def format_complex(self, val: Union[sympy.Basic, int, float, 'cirq.TParamValComplex']) -> str:
+    def format_complex(self, val: sympy.Basic | int | float | cirq.TParamValComplex) -> str:
         if isinstance(val, sympy.Basic):
             return str(val)
         c = complex(val)
@@ -263,7 +282,7 @@ class CircuitDiagramInfoArgs:
         imag_str = '' if abs_imag == 1 else self.format_real(abs_imag)
         return f'{self.format_real(c.real)}{joiner}{imag_str}i'
 
-    def format_radians(self, radians: Union[sympy.Basic, int, float]) -> str:
+    def format_radians(self, radians: sympy.Basic | int | float) -> str:
         """Returns angle in radians as a human-readable string."""
         if protocols.is_parameterized(radians):
             return str(radians)
@@ -288,6 +307,7 @@ class CircuitDiagramInfoArgs:
             use_unicode_characters=self.use_unicode_characters,
             precision=self.precision,
             label_map=self.label_map,
+            include_tags=self.include_tags,
             transpose=self.transpose,
         )
 
@@ -314,7 +334,7 @@ class SupportsCircuitDiagramInfo(Protocol):
     @doc_private
     def _circuit_diagram_info_(
         self, args: CircuitDiagramInfoArgs
-    ) -> Union[str, Iterable[str], CircuitDiagramInfo]:
+    ) -> str | Iterable[str] | CircuitDiagramInfo:
         """Describes how to draw an operation in a circuit diagram.
 
         This method is used by the global `cirq.diagram_info` method. If this
@@ -336,10 +356,10 @@ RaiseTypeErrorIfNotProvided = CircuitDiagramInfo(())
 
 
 def _op_info_with_fallback(
-    op: 'cirq.Operation', args: 'cirq.CircuitDiagramInfoArgs'
-) -> 'cirq.CircuitDiagramInfo':
+    op: cirq.Operation, args: cirq.CircuitDiagramInfoArgs
+) -> cirq.CircuitDiagramInfo:
     info = protocols.circuit_diagram_info(op, args, None)
-    rows: List[LabelEntity] = list(op.qubits)
+    rows: list[LabelEntity] = list(op.qubits)
     if args.label_map is not None:
         rows += protocols.measurement_keys_touched(op) & args.label_map.keys()
     if info is not None and info.wire_symbols:
@@ -366,28 +386,27 @@ def _op_info_with_fallback(
     return protocols.CircuitDiagramInfo(wire_symbols=symbols)
 
 
-# pylint: disable=function-redefined
 @overload
 def circuit_diagram_info(
-    val: Any, args: Optional[CircuitDiagramInfoArgs] = None
+    val: Any, args: CircuitDiagramInfoArgs | None = None
 ) -> CircuitDiagramInfo:
     pass
 
 
 @overload
 def circuit_diagram_info(
-    val: Any, args: Optional[CircuitDiagramInfoArgs], default: TDefault
-) -> Union[CircuitDiagramInfo, TDefault]:
+    val: Any, args: CircuitDiagramInfoArgs | None, default: TDefault
+) -> CircuitDiagramInfo | TDefault:
     pass
 
 
 @overload
-def circuit_diagram_info(val: Any, *, default: TDefault) -> Union[CircuitDiagramInfo, TDefault]:
+def circuit_diagram_info(val: Any, *, default: TDefault) -> CircuitDiagramInfo | TDefault:
     pass
 
 
 def circuit_diagram_info(
-    val: Any, args: Optional[CircuitDiagramInfoArgs] = None, default=RaiseTypeErrorIfNotProvided
+    val: Any, args: CircuitDiagramInfoArgs | None = None, default=RaiseTypeErrorIfNotProvided
 ):
     """Requests information on drawing an operation in a circuit diagram.
 
@@ -438,6 +457,3 @@ def circuit_diagram_info(
         f"object of type '{type(val)}' does have a _circuit_diagram_info_ "
         "method, but it returned NotImplemented."
     )
-
-
-# pylint: enable=function-redefined