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

cirq/circuits/frozen_circuit.py

@@ -18,25 +18,15 @@ from __future__ import annotations
 
 from functools import cached_property
 from types import NotImplementedType
-from typing import (
-    AbstractSet,
-    FrozenSet,
-    Hashable,
-    Iterable,
-    Iterator,
-    Sequence,
-    Tuple,
-    TYPE_CHECKING,
-    Union,
-)
-
-import numpy as np
+from typing import AbstractSet, Hashable, Iterable, Iterator, Sequence, TYPE_CHECKING
 
 from cirq import _compat, protocols
 from cirq.circuits import AbstractCircuit, Alignment, Circuit
 from cirq.circuits.insert_strategy import InsertStrategy
 
 if TYPE_CHECKING:
+    import numpy as np
+
     import cirq
 
 
@@ -75,9 +65,12 @@ class FrozenCircuit(AbstractCircuit, protocols.SerializableByKey):
         self._tags = tuple(tags)
 
     @classmethod
-    def _from_moments(cls, moments: Iterable[cirq.Moment]) -> FrozenCircuit:
+    def _from_moments(
+        cls, moments: Iterable[cirq.Moment], tags: Sequence[Hashable]
+    ) -> FrozenCircuit:
         new_circuit = FrozenCircuit()
         new_circuit._moments = tuple(moments)
+        new_circuit._tags = tuple(tags)
         return new_circuit
 
     @property
@@ -88,17 +81,16 @@ class FrozenCircuit(AbstractCircuit, protocols.SerializableByKey):
         return self
 
     def unfreeze(self, copy: bool = True) -> cirq.Circuit:
-        return Circuit._from_moments(self._moments)
+        return Circuit._from_moments(self._moments, tags=self.tags)
 
     @property
-    def tags(self) -> Tuple[Hashable, ...]:
+    def tags(self) -> tuple[Hashable, ...]:
         """Returns a tuple of the Circuit's tags."""
         return self._tags
 
     @cached_property
     def untagged(self) -> cirq.FrozenCircuit:
-        """Returns the underlying FrozenCircuit without any tags."""
-        return self._from_moments(self._moments) if self.tags else self
+        return super().untagged
 
     def with_tags(self, *new_tags: Hashable) -> cirq.FrozenCircuit:
         """Creates a new tagged `FrozenCircuit` with `self.tags` and `new_tags` combined."""
@@ -113,13 +105,6 @@ class FrozenCircuit(AbstractCircuit, protocols.SerializableByKey):
         # Explicitly cached for performance
         return hash((self.moments, self.tags))
 
-    def __eq__(self, other):
-        super_eq = super().__eq__(other)
-        if super_eq is not True:
-            return super_eq
-        other_tags = other.tags if isinstance(other, FrozenCircuit) else ()
-        return self.tags == other_tags
-
     def __getstate__(self):
         # Don't save hash when pickling; see #3777.
         state = self.__dict__
@@ -134,7 +119,7 @@ class FrozenCircuit(AbstractCircuit, protocols.SerializableByKey):
         return len(self.all_qubits())
 
     @_compat.cached_method
-    def _qid_shape_(self) -> Tuple[int, ...]:
+    def _qid_shape_(self) -> tuple[int, ...]:
         return super()._qid_shape_()
 
     @_compat.cached_method
@@ -142,7 +127,7 @@ class FrozenCircuit(AbstractCircuit, protocols.SerializableByKey):
         return super()._has_unitary_()
 
     @_compat.cached_method
-    def _unitary_(self) -> Union[np.ndarray, NotImplementedType]:
+    def _unitary_(self) -> np.ndarray | NotImplementedType:
         return super()._unitary_()
 
     @_compat.cached_method
@@ -150,11 +135,11 @@ class FrozenCircuit(AbstractCircuit, protocols.SerializableByKey):
         return protocols.is_measurement(self.unfreeze())
 
     @_compat.cached_method
-    def all_qubits(self) -> FrozenSet[cirq.Qid]:
+    def all_qubits(self) -> frozenset[cirq.Qid]:
         return super().all_qubits()
 
     @cached_property
-    def _all_operations(self) -> Tuple[cirq.Operation, ...]:
+    def _all_operations(self) -> tuple[cirq.Operation, ...]:
         return tuple(super().all_operations())
 
     def all_operations(self) -> Iterator[cirq.Operation]:
@@ -164,14 +149,14 @@ class FrozenCircuit(AbstractCircuit, protocols.SerializableByKey):
         return self._is_measurement_()
 
     @_compat.cached_method
-    def all_measurement_key_objs(self) -> FrozenSet[cirq.MeasurementKey]:
+    def all_measurement_key_objs(self) -> frozenset[cirq.MeasurementKey]:
         return super().all_measurement_key_objs()
 
-    def _measurement_key_objs_(self) -> FrozenSet[cirq.MeasurementKey]:
+    def _measurement_key_objs_(self) -> frozenset[cirq.MeasurementKey]:
         return self.all_measurement_key_objs()
 
     @_compat.cached_method
-    def _control_keys_(self) -> FrozenSet[cirq.MeasurementKey]:
+    def _control_keys_(self) -> frozenset[cirq.MeasurementKey]:
         return super()._control_keys_()
 
     @_compat.cached_method
@@ -179,30 +164,18 @@ class FrozenCircuit(AbstractCircuit, protocols.SerializableByKey):
         return super().are_all_measurements_terminal()
 
     @_compat.cached_method
-    def all_measurement_key_names(self) -> FrozenSet[str]:
+    def all_measurement_key_names(self) -> frozenset[str]:
         return frozenset(str(key) for key in self.all_measurement_key_objs())
 
     @_compat.cached_method
     def _is_parameterized_(self) -> bool:
-        return super()._is_parameterized_() or any(
-            protocols.is_parameterized(tag) for tag in self.tags
-        )
+        return super()._is_parameterized_()
 
     @_compat.cached_method
     def _parameter_names_(self) -> AbstractSet[str]:
-        tag_params = {name for tag in self.tags for name in protocols.parameter_names(tag)}
-        return super()._parameter_names_() | tag_params
-
-    def _resolve_parameters_(
-        self, resolver: cirq.ParamResolver, recursive: bool
-    ) -> cirq.FrozenCircuit:
-        resolved_circuit = super()._resolve_parameters_(resolver, recursive)
-        resolved_tags = [
-            protocols.resolve_parameters(tag, resolver, recursive) for tag in self.tags
-        ]
-        return resolved_circuit.with_tags(*resolved_tags)
+        return super()._parameter_names_()
 
-    def _measurement_key_names_(self) -> FrozenSet[str]:
+    def _measurement_key_names_(self) -> frozenset[str]:
         return self.all_measurement_key_names()
 
     def __add__(self, other) -> cirq.FrozenCircuit:
@@ -227,29 +200,19 @@ class FrozenCircuit(AbstractCircuit, protocols.SerializableByKey):
         except:
             return NotImplemented
 
-    def _repr_args(self) -> str:
-        moments_repr = super()._repr_args()
-        tag_repr = ','.join(_compat.proper_repr(t) for t in self._tags)
-        return f'{moments_repr}, tags=[{tag_repr}]' if self.tags else moments_repr
-
-    def _json_dict_(self):
-        attribute_names = ['moments', 'tags'] if self.tags else ['moments']
-        ret = protocols.obj_to_dict_helper(self, attribute_names)
-        return ret
-
     @classmethod
     def _from_json_dict_(cls, moments, *, tags=(), **kwargs):
         return cls(moments, strategy=InsertStrategy.EARLIEST, tags=tags)
 
     def concat_ragged(
-        *circuits: cirq.AbstractCircuit, align: Union[cirq.Alignment, str] = Alignment.LEFT
+        *circuits: cirq.AbstractCircuit, align: cirq.Alignment | str = Alignment.LEFT
     ) -> cirq.FrozenCircuit:
         return AbstractCircuit.concat_ragged(*circuits, align=align).freeze()
 
     concat_ragged.__doc__ = AbstractCircuit.concat_ragged.__doc__
 
     def zip(
-        *circuits: cirq.AbstractCircuit, align: Union[cirq.Alignment, str] = Alignment.LEFT
+        *circuits: cirq.AbstractCircuit, align: cirq.Alignment | str = Alignment.LEFT
     ) -> cirq.FrozenCircuit:
         return AbstractCircuit.zip(*circuits, align=align).freeze()
 

cirq/circuits/frozen_circuit_test.py

@@ -16,13 +16,15 @@
 Behavior shared with Circuit is tested with parameters in circuit_test.py.
 """
 
+from __future__ import annotations
+
 import pytest
 import sympy
 
 import cirq
 
 
-def test_from_moments():
+def test_from_moments() -> None:
     a, b, c, d = cirq.LineQubit.range(4)
     moment = cirq.Moment(cirq.Z(a), cirq.Z(b))
     subcircuit = cirq.FrozenCircuit.from_moments(cirq.X(c), cirq.Y(d))
@@ -49,10 +51,11 @@ def test_from_moments():
         cirq.Moment(cirq.measure(a, b, key='ab'), cirq.measure(c, d, key='cd')),
     )
     assert circuit[0] is moment
+    assert isinstance(circuit[1].operations[0], cirq.CircuitOperation)
     assert circuit[1].operations[0].circuit is subcircuit
 
 
-def test_freeze_and_unfreeze():
+def test_freeze_and_unfreeze() -> None:
     a, b = cirq.LineQubit.range(2)
     c = cirq.Circuit(cirq.X(a), cirq.H(b))
 
@@ -78,7 +81,7 @@ def test_freeze_and_unfreeze():
     assert fcc is not f
 
 
-def test_immutable():
+def test_immutable() -> None:
     q = cirq.LineQubit(0)
     c = cirq.FrozenCircuit(cirq.X(q), cirq.H(q))
 
@@ -87,21 +90,21 @@ def test_immutable():
         AttributeError,
         match="(can't set attribute)|(property 'moments' of 'FrozenCircuit' object has no setter)",
     ):
-        c.moments = (cirq.Moment(cirq.H(q)), cirq.Moment(cirq.X(q)))
+        c.moments = (cirq.Moment(cirq.H(q)), cirq.Moment(cirq.X(q)))  # type: ignore[misc]
 
 
-def test_tagged_circuits():
+def test_tagged_circuits() -> None:
     q = cirq.LineQubit(0)
     ops = [cirq.X(q), cirq.H(q)]
-    tags = [sympy.Symbol("a"), "b"]
+    tags = (sympy.Symbol("a"), "b")
     circuit = cirq.Circuit(ops)
     frozen_circuit = cirq.FrozenCircuit(ops)
     tagged_circuit = cirq.FrozenCircuit(ops, tags=tags)
     # Test equality
-    assert tagged_circuit.tags == tuple(tags)
+    assert tagged_circuit.tags == tags
     assert circuit == frozen_circuit != tagged_circuit
+    assert not cirq.approx_eq(frozen_circuit, tagged_circuit)
     assert cirq.approx_eq(circuit, frozen_circuit)
-    assert cirq.approx_eq(frozen_circuit, tagged_circuit)
     # Test hash
     assert hash(frozen_circuit) != hash(tagged_circuit)
     # Test _repr_ and _json_ round trips.
@@ -113,9 +116,29 @@ def test_tagged_circuits():
     assert tagged_circuit.with_tags("c") == cirq.FrozenCircuit(ops, tags=[*tags, "c"])
     assert tagged_circuit.untagged == frozen_circuit
     assert frozen_circuit.untagged is frozen_circuit
+    assert tagged_circuit.unfreeze(copy=True).tags == tags
+    assert tagged_circuit.unfreeze(copy=False).tags == tags
     # Test parameterized protocols
     assert cirq.is_parameterized(frozen_circuit) is False
     assert cirq.is_parameterized(tagged_circuit) is True
     assert cirq.parameter_names(tagged_circuit) == {"a"}
+    assert cirq.resolve_parameters(tagged_circuit, {"a": 1}).tags == (1, "b")
     # Tags are not propagated to diagrams yet.
     assert str(frozen_circuit) == str(tagged_circuit)
+    # Test tags are preserved through operations
+    assert (tagged_circuit + circuit).tags == tags
+    assert (circuit + tagged_circuit).tags == ()  # We only preserve the tags for the first one
+    assert (2 * tagged_circuit).tags == tags
+    assert (tagged_circuit * 2).tags == tags
+    assert (tagged_circuit**-1).tags == tags
+    for c in tagged_circuit.factorize():
+        assert tagged_circuit.tags == tags
+
+    q2 = cirq.LineQubit(1)
+    circuit2 = cirq.Circuit(cirq.X(q2), cirq.H(q2))
+    assert tagged_circuit.zip(circuit2).tags == tags
+    assert circuit2.zip(tagged_circuit).tags == ()  # We only preserve the tags for the first one
+    assert tagged_circuit.concat_ragged(circuit2).tags == tags
+    assert (
+        circuit2.concat_ragged(tagged_circuit).tags == ()
+    )  # We only preserve the tags for the first one

cirq/circuits/insert_strategy.py

@@ -14,16 +14,18 @@
 
 """Hard-coded options for adding multiple operations to a circuit."""
 
+from __future__ import annotations
+
 
 class InsertStrategy:
     """Indicates preferences on how to add multiple operations to a circuit."""
 
-    NEW: 'InsertStrategy'
-    NEW_THEN_INLINE: 'InsertStrategy'
-    INLINE: 'InsertStrategy'
-    EARLIEST: 'InsertStrategy'
+    NEW: InsertStrategy
+    NEW_THEN_INLINE: InsertStrategy
+    INLINE: InsertStrategy
+    EARLIEST: InsertStrategy
 
-    def __new__(cls, name: str, doc: str) -> 'InsertStrategy':
+    def __new__(cls, name: str, doc: str) -> InsertStrategy:
         inst = getattr(cls, name, None)
         if not inst or not isinstance(inst, cls):
             inst = super().__new__(cls)

cirq/circuits/insert_strategy_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 pickle
 
 import pytest
@@ -19,7 +21,7 @@ import pytest
 import cirq
 
 
-def test_repr():
+def test_repr() -> None:
     assert repr(cirq.InsertStrategy.NEW) == 'cirq.InsertStrategy.NEW'
     assert str(cirq.InsertStrategy.NEW) == 'NEW'
 
@@ -34,6 +36,6 @@ def test_repr():
     ],
     ids=lambda strategy: strategy.name,
 )
-def test_identity_after_pickling(strategy: cirq.InsertStrategy):
+def test_identity_after_pickling(strategy: cirq.InsertStrategy) -> None:
     unpickled_strategy = pickle.loads(pickle.dumps(strategy))
     assert unpickled_strategy is strategy

cirq/circuits/moment.py

@@ -24,18 +24,13 @@ from typing import (
     Any,
     Callable,
     cast,
-    Dict,
-    FrozenSet,
+    Hashable,
     Iterable,
     Iterator,
-    List,
     Mapping,
-    Optional,
     overload,
     Sequence,
-    Tuple,
     TYPE_CHECKING,
-    Union,
 )
 
 import numpy as np
@@ -51,13 +46,12 @@ if TYPE_CHECKING:
 
 # Lazy imports to break circular dependencies.
 circuit = LazyLoader("circuit", globals(), "cirq.circuits.circuit")
-op_tree = LazyLoader("op_tree", globals(), "cirq.ops.op_tree")
 text_diagram_drawer = LazyLoader(
     "text_diagram_drawer", globals(), "cirq.circuits.text_diagram_drawer"
 )
 
 
-def _default_breakdown(qid: cirq.Qid) -> Tuple[Any, Any]:
+def _default_breakdown(qid: cirq.Qid) -> tuple[Any, Any]:
     # Attempt to convert into a position on the complex plane.
     try:
         plane_pos = complex(qid)  # type: ignore
@@ -84,7 +78,12 @@ class Moment:
             are no such operations, returns an empty Moment.
     """
 
-    def __init__(self, *contents: cirq.OP_TREE, _flatten_contents: bool = True) -> None:
+    def __init__(
+        self,
+        *contents: cirq.OP_TREE,
+        _flatten_contents: bool = True,
+        tags: tuple[Hashable, ...] = (),
+    ) -> None:
         """Constructs a moment with the given operations.
 
         Args:
@@ -95,6 +94,12 @@ class Moment:
                 we skip flattening and assume that contents already consists
                 of individual operations. This is used internally by helper
                 methods to avoid unnecessary validation.
+            tags:  Optional tags to denote specific Moment objects with meta-data.
+                These are a tuple of any Hashable object.  Typically, a class
+                will be passed.  Tags apply only to this specific set of operations
+                and will be lost on any transformation of the
+                Moment.  For instance, if operations are added to the Moment, tags
+                will be dropped unless explicitly added back in by the user.
 
         Raises:
             ValueError: A qubit appears more than once.
@@ -102,12 +107,12 @@ class Moment:
         self._operations = (
             tuple(op_tree.flatten_to_ops(contents))
             if _flatten_contents
-            else cast(Tuple['cirq.Operation'], contents)
+            else cast(tuple['cirq.Operation'], contents)
         )
-        self._sorted_operations: Optional[Tuple[cirq.Operation, ...]] = None
+        self._sorted_operations: tuple[cirq.Operation, ...] | None = None
 
         # An internal dictionary to support efficient operation access by qubit.
-        self._qubit_to_op: Dict[cirq.Qid, cirq.Operation] = {}
+        self._qubit_to_op: dict[cirq.Qid, cirq.Operation] = {}
         for op in self.operations:
             for q in op.qubits:
                 # Check that operations don't overlap.
@@ -115,11 +120,12 @@ class Moment:
                     raise ValueError(f'Overlapping operations: {self.operations}')
                 self._qubit_to_op[q] = op
 
-        self._measurement_key_objs: Optional[FrozenSet[cirq.MeasurementKey]] = None
-        self._control_keys: Optional[FrozenSet[cirq.MeasurementKey]] = None
+        self._measurement_key_objs: frozenset[cirq.MeasurementKey] | None = None
+        self._control_keys: frozenset[cirq.MeasurementKey] | None = None
+        self._tags = tags
 
     @classmethod
-    def from_ops(cls, *ops: cirq.Operation) -> cirq.Moment:
+    def from_ops(cls, *ops: cirq.Operation, tags: tuple[Hashable, ...] = ()) -> cirq.Moment:
         """Construct a Moment from the given operations.
 
         This avoids calling `flatten_to_ops` in the moment constructor, which
@@ -129,17 +135,48 @@ class Moment:
 
         Args:
             *ops: Operations to include in the Moment.
+            tags:  Optional tags to denote specific Moment objects with meta-data.
+                These are a tuple of any Hashable object.  Tags will be dropped if
+                the operations in the Moment are modified or transformed.
         """
-        return cls(*ops, _flatten_contents=False)
+        return cls(*ops, _flatten_contents=False, tags=tags)
 
     @property
-    def operations(self) -> Tuple[cirq.Operation, ...]:
+    def operations(self) -> tuple[cirq.Operation, ...]:
         return self._operations
 
     @cached_property
-    def qubits(self) -> FrozenSet[cirq.Qid]:
+    def qubits(self) -> frozenset[cirq.Qid]:
         return frozenset(self._qubit_to_op)
 
+    @property
+    def tags(self) -> tuple[Hashable, ...]:
+        """Returns a tuple of the operation's tags."""
+        return self._tags
+
+    def with_tags(self, *new_tags: Hashable) -> cirq.Moment:
+        """Creates a new Moment with the current ops and the specified tags.
+
+        If the moment already has tags, this will add the new_tags to the
+        preexisting tags.
+
+        This method can be used to attach meta-data to moments
+        without affecting their functionality.  The intended usage is to
+        attach classes intended for this purpose or strings to mark operations
+        for specific usage that will be recognized by consumers.
+
+        Tags can be a list of any type of object that is useful to identify
+        this operation as long as the type is hashable.  If you wish the
+        resulting operation to be eventually serialized into JSON, you should
+        also restrict the operation to be JSON serializable.
+
+        Please note that tags should be instantiated if classes are
+        used.  Raw types are not allowed.
+        """
+        if not new_tags:
+            return self
+        return Moment(*self._operations, _flatten_contents=False, tags=(*self._tags, *new_tags))
+
     def operates_on_single_qubit(self, qubit: cirq.Qid) -> bool:
         """Determines if the moment has operations touching the given qubit.
         Args:
@@ -160,7 +197,7 @@ class Moment:
         """
         return not self._qubit_to_op.keys().isdisjoint(qubits)
 
-    def operation_at(self, qubit: raw_types.Qid) -> Optional[cirq.Operation]:
+    def operation_at(self, qubit: raw_types.Qid) -> cirq.Operation | None:
         """Returns the operation on a certain qubit for the moment.
 
         Args:
@@ -177,6 +214,8 @@ class Moment:
     def with_operation(self, operation: cirq.Operation) -> cirq.Moment:
         """Returns an equal moment, but with the given op added.
 
+        Any tags on the Moment will be dropped.
+
         Args:
             operation: The operation to append.
 
@@ -205,6 +244,9 @@ class Moment:
     def with_operations(self, *contents: cirq.OP_TREE) -> cirq.Moment:
         """Returns a new moment with the given contents added.
 
+        Any tags on the original Moment object are dropped if the Moment
+        is changed.
+
         Args:
             *contents: New operations to add to this moment.
 
@@ -242,6 +284,9 @@ class Moment:
     def without_operations_touching(self, qubits: Iterable[cirq.Qid]) -> cirq.Moment:
         """Returns an equal moment, but without ops on the given qubits.
 
+        Any tags on the original Moment object are dropped if the Moment
+        is changed.
+
         Args:
             qubits: Operations that touch these will be removed.
 
@@ -267,7 +312,7 @@ class Moment:
 
     def _resolve_parameters_(self, resolver: cirq.ParamResolver, recursive: bool) -> cirq.Moment:
         changed = False
-        resolved_ops: List[cirq.Operation] = []
+        resolved_ops: list[cirq.Operation] = []
         for op in self:
             resolved_op = protocols.resolve_parameters(op, resolver, recursive)
             changed = (
@@ -291,35 +336,35 @@ class Moment:
         )
 
     @_compat.cached_method()
-    def _measurement_key_names_(self) -> FrozenSet[str]:
+    def _measurement_key_names_(self) -> frozenset[str]:
         return frozenset(str(key) for key in self._measurement_key_objs_())
 
-    def _measurement_key_objs_(self) -> FrozenSet[cirq.MeasurementKey]:
+    def _measurement_key_objs_(self) -> frozenset[cirq.MeasurementKey]:
         if self._measurement_key_objs is None:
             self._measurement_key_objs = frozenset(
                 key for op in self.operations for key in protocols.measurement_key_objs(op)
             )
         return self._measurement_key_objs
 
-    def _control_keys_(self) -> FrozenSet[cirq.MeasurementKey]:
+    def _control_keys_(self) -> frozenset[cirq.MeasurementKey]:
         if self._control_keys is None:
             self._control_keys = frozenset(
                 k for op in self.operations for k in protocols.control_keys(op)
             )
         return self._control_keys
 
-    def _sorted_operations_(self) -> Tuple[cirq.Operation, ...]:
+    def _sorted_operations_(self) -> tuple[cirq.Operation, ...]:
         if self._sorted_operations is None:
             self._sorted_operations = tuple(sorted(self._operations, key=lambda op: op.qubits))
         return self._sorted_operations
 
-    def _with_key_path_(self, path: Tuple[str, ...]):
+    def _with_key_path_(self, path: tuple[str, ...]):
         return Moment(
             protocols.with_key_path(op, path) if protocols.is_measurement(op) else op
             for op in self.operations
         )
 
-    def _with_key_path_prefix_(self, prefix: Tuple[str, ...]):
+    def _with_key_path_prefix_(self, prefix: tuple[str, ...]):
         return Moment(
             (
                 protocols.with_key_path_prefix(op, prefix)
@@ -330,7 +375,7 @@ class Moment:
         )
 
     def _with_rescoped_keys_(
-        self, path: Tuple[str, ...], bindable_keys: FrozenSet[cirq.MeasurementKey]
+        self, path: tuple[str, ...], bindable_keys: frozenset[cirq.MeasurementKey]
     ):
         return Moment(
             protocols.with_rescoped_keys(op, path, bindable_keys) for op in self.operations
@@ -364,7 +409,7 @@ class Moment:
     def __hash__(self):
         return hash((Moment, self._sorted_operations_()))
 
-    def __getstate__(self) -> Dict[str, Any]:
+    def __getstate__(self) -> dict[str, Any]:
         # clear cached hash value when pickling, see #6674
         state = self.__dict__
         hash_attr = _compat._method_cache_name(self.__hash__)
@@ -407,7 +452,7 @@ class Moment:
         return self._operations
 
     def transform_qubits(
-        self, qubit_map: Union[Dict[cirq.Qid, cirq.Qid], Callable[[cirq.Qid], cirq.Qid]]
+        self, qubit_map: dict[cirq.Qid, cirq.Qid] | Callable[[cirq.Qid], cirq.Qid]
     ) -> Self:
         """Returns the same moment, but with different qubits.
 
@@ -516,12 +561,14 @@ class Moment:
             return NotImplemented
         return qis.kraus_to_superoperator(self._kraus_())
 
-    def _json_dict_(self) -> Dict[str, Any]:
-        return protocols.obj_to_dict_helper(self, ['operations'])
+    def _json_dict_(self) -> dict[str, Any]:
+        # For backwards compatibility, only output tags if they exist.
+        args = ['operations', 'tags'] if self._tags else ['operations']
+        return protocols.obj_to_dict_helper(self, args)
 
     @classmethod
-    def _from_json_dict_(cls, operations, **kwargs):
-        return cls.from_ops(*operations)
+    def _from_json_dict_(cls, operations, tags=(), **kwargs):
+        return cls(*operations, tags=tags)
 
     def __add__(self, other: cirq.OP_TREE) -> cirq.Moment:
         if isinstance(other, circuit.AbstractCircuit):
@@ -544,7 +591,6 @@ class Moment:
             )
         return Moment(new_ops)
 
-    # pylint: disable=function-redefined
     @overload
     def __getitem__(self, key: raw_types.Qid) -> cirq.Operation:
         pass
@@ -569,11 +615,11 @@ class Moment:
     def to_text_diagram(
         self: cirq.Moment,
         *,
-        xy_breakdown_func: Callable[[cirq.Qid], Tuple[Any, Any]] = _default_breakdown,
+        xy_breakdown_func: Callable[[cirq.Qid], tuple[Any, Any]] = _default_breakdown,
         extra_qubits: Iterable[cirq.Qid] = (),
         use_unicode_characters: bool = True,
-        precision: Optional[int] = None,
-        include_tags: bool = True,
+        precision: int | None = None,
+        include_tags: bool | Iterable[type] = True,
     ) -> str:
         """Create a text diagram for the moment.
 
@@ -591,8 +637,10 @@ class Moment:
             precision: How precise numbers, such as angles, should be. Use None
                 for infinite precision, or an integer for a certain number of
                 digits of precision.
-            include_tags: Whether or not to include operation tags in the
-                diagram.
+            include_tags: Controls which tags attached to operations are
+                included. ``True`` includes all tags, ``False`` includes none,
+                or a collection of tag classes may be specified to include only
+                those tags.
 
         Returns:
             The text diagram rendered into text.
@@ -657,7 +705,7 @@ class Moment:
 
         return diagram.render()
 
-    def _commutes_(self, other: Any, *, atol: float = 1e-8) -> Union[bool, NotImplementedType]:
+    def _commutes_(self, other: Any, *, atol: float = 1e-8) -> bool | NotImplementedType:
         """Determines whether Moment commutes with the other Moment or Operation.
 
         Args: