Repository: quantumlib/Cirq

Issue:
Repeated CircuitOperation with measurements
**Is your feature request related to a use case or problem? Please describe.**
Currently, CircuitOperation does not allow repeating the sub-circuit if there is a measurement in it. This feature is essential for flow-control (#3234) through repeated/conditional substructures.

**Describe the solution you'd like**
Use repetition identifiers as measurement key prefixes to uniquely identify each measurement
RFC: https://tinyurl.com/measurement-repeated-circuitop

**What is the urgency from your perspective for this issue? Is it blocking important work?**
P2 - we should do it in the next couple of quarters 


Discussion:

Resolving patch:
diff --git a/cirq/circuits/circuit_operation.py b/cirq/circuits/circuit_operation.py
--- a/cirq/circuits/circuit_operation.py
+++ b/cirq/circuits/circuit_operation.py
@@ -18,7 +18,7 @@
 component operations in order, including any nested CircuitOperations.
 """
 
-from typing import TYPE_CHECKING, AbstractSet, Callable, Dict, Optional, Tuple, Union
+from typing import TYPE_CHECKING, AbstractSet, Callable, Dict, List, Optional, Tuple, Union
 
 import dataclasses
 import numpy as np
@@ -33,6 +33,22 @@
 INT_TYPE = Union[int, np.integer]
 
 
+def default_repetition_ids(repetitions: int) -> Optional[List[str]]:
+    if abs(repetitions) > 1:
+        return [str(i) for i in range(abs(repetitions))]
+    return None
+
+
+def cartesian_product_of_string_lists(list1: Optional[List[str]], list2: Optional[List[str]]):
+    if list1 is None and list2 is None:
+        return None  # coverage: ignore
+    if list1 is None:
+        return list2  # coverage: ignore
+    if list2 is None:
+        return list1
+    return [f'{first}-{second}' for first in list1 for second in list2]
+
+
 @dataclasses.dataclass(frozen=True)
 class CircuitOperation(ops.Operation):
     """An operation that encapsulates a circuit.
@@ -47,6 +63,10 @@ class CircuitOperation(ops.Operation):
         qubit_map: Remappings for qubits in the circuit.
         measurement_key_map: Remappings for measurement keys in the circuit.
         param_resolver: Resolved values for parameters in the circuit.
+        repetition_ids: List of identifiers for each repetition of the
+            CircuitOperation. If populated, the length should be equal to the
+            repetitions. If not populated and abs(`repetitions`) > 1, it is
+            initialized to strings for numbers in `range(repetitions)`.
     """
 
     _hash: Optional[int] = dataclasses.field(default=None, init=False)
@@ -56,10 +76,28 @@ class CircuitOperation(ops.Operation):
     qubit_map: Dict['cirq.Qid', 'cirq.Qid'] = dataclasses.field(default_factory=dict)
     measurement_key_map: Dict[str, str] = dataclasses.field(default_factory=dict)
     param_resolver: study.ParamResolver = study.ParamResolver()
+    repetition_ids: Optional[List[str]] = dataclasses.field(default=None)
 
     def __post_init__(self):
         if not isinstance(self.circuit, circuits.FrozenCircuit):
             raise TypeError(f'Expected circuit of type FrozenCircuit, got: {type(self.circuit)!r}')
+
+        # Ensure that the circuit is invertible if the repetitions are negative.
+        if self.repetitions < 0:
+            try:
+                protocols.inverse(self.circuit.unfreeze())
+            except TypeError:
+                raise ValueError(f'repetitions are negative but the circuit is not invertible')
+
+        # Initialize repetition_ids to default, if unspecified. Else, validate their length.
+        loop_size = abs(self.repetitions)
+        if not self.repetition_ids:
+            object.__setattr__(self, 'repetition_ids', self._default_repetition_ids())
+        elif len(self.repetition_ids) != loop_size:
+            raise ValueError(
+                f'Expected repetition_ids to be a list of length {loop_size}, '
+                f'got: {self.repetition_ids}'
+            )
         # Ensure that param_resolver is converted to an actual ParamResolver.
         object.__setattr__(self, 'param_resolver', study.ParamResolver(self.param_resolver))
 
@@ -83,6 +121,7 @@ def __eq__(self, other) -> bool:
             and self.measurement_key_map == other.measurement_key_map
             and self.param_resolver == other.param_resolver
             and self.repetitions == other.repetitions
+            and self.repetition_ids == other.repetition_ids
         )
 
     # Methods for getting post-mapping properties of the contained circuit.
@@ -93,6 +132,9 @@ def qubits(self) -> Tuple['cirq.Qid', ...]:
         ordered_qubits = ops.QubitOrder.DEFAULT.order_for(self.circuit.all_qubits())
         return tuple(self.qubit_map.get(q, q) for q in ordered_qubits)
 
+    def _default_repetition_ids(self) -> Optional[List[str]]:
+        return default_repetition_ids(self.repetitions)
+
     def _qid_shape_(self) -> Tuple[int, ...]:
         return tuple(q.dimension for q in self.qubits)
 
@@ -117,8 +159,39 @@ def _decompose_(self) -> 'cirq.OP_TREE':
             result = result ** -1
         result = protocols.with_measurement_key_mapping(result, self.measurement_key_map)
         result = protocols.resolve_parameters(result, self.param_resolver, recursive=False)
-
-        return list(result.all_operations()) * abs(self.repetitions)
+        # repetition_ids don't need to be taken into account if the circuit has no measurements
+        # or if repetition_ids are unset.
+        if self.repetition_ids is None or not protocols.is_measurement(result):
+            return list(result.all_operations()) * abs(self.repetitions)
+        # If it's a measurement circuit with repetitions/repetition_ids, prefix the repetition_ids
+        # to measurements. Details at https://tinyurl.com/measurement-repeated-circuitop.
+        ops = []  # type: List[cirq.Operation]
+        for parent_id in self.repetition_ids:
+            for op in result.all_operations():
+                if isinstance(op, CircuitOperation):
+                    # For a CircuitOperation, prefix the current repetition_id to the children
+                    # repetition_ids.
+                    ops.append(
+                        op.with_repetition_ids(
+                            # If `op.repetition_ids` is None, this will return `[parent_id]`.
+                            cartesian_product_of_string_lists([parent_id], op.repetition_ids)
+                        )
+                    )
+                elif protocols.is_measurement(op):
+                    # For a non-CircuitOperation measurement, prefix the current repetition_id
+                    # to the children measurement keys. Implemented by creating a mapping and
+                    # using the with_measurement_key_mapping protocol.
+                    ops.append(
+                        protocols.with_measurement_key_mapping(
+                            op,
+                            key_map={
+                                key: f'{parent_id}-{key}' for key in protocols.measurement_keys(op)
+                            },
+                        )
+                    )
+                else:
+                    ops.append(op)
+        return ops
 
     # Methods for string representation of the operation.
 
@@ -132,6 +205,9 @@ def __repr__(self):
             args += f'measurement_key_map={proper_repr(self.measurement_key_map)},\n'
         if self.param_resolver:
             args += f'param_resolver={proper_repr(self.param_resolver)},\n'
+        if self.repetition_ids != self._default_repetition_ids():
+            # Default repetition_ids need not be specified.
+            args += f'repetition_ids={proper_repr(self.repetition_ids)},\n'
         indented_args = args.replace('\n', '\n    ')
         return f'cirq.CircuitOperation({indented_args[:-4]})'
 
@@ -155,7 +231,11 @@ def dict_str(d: Dict) -> str:
             args.append(f'key_map={dict_str(self.measurement_key_map)}')
         if self.param_resolver:
             args.append(f'params={self.param_resolver.param_dict}')
-        if self.repetitions != 1:
+        if self.repetition_ids != self._default_repetition_ids():
+            # Default repetition_ids need not be specified.
+            args.append(f'repetition_ids={self.repetition_ids}')
+        elif self.repetitions != 1:
+            # Only add loops if we haven't added repetition_ids.
             args.append(f'loops={self.repetitions}')
         if not args:
             return f'{header}\n{circuit_msg}'
@@ -173,6 +253,7 @@ def __hash__(self):
                         frozenset(self.qubit_map.items()),
                         frozenset(self.measurement_key_map.items()),
                         self.param_resolver,
+                        tuple([] if self.repetition_ids is None else self.repetition_ids),
                     )
                 ),
             )
@@ -188,53 +269,95 @@ def _json_dict_(self):
             'qubit_map': sorted(self.qubit_map.items()),
             'measurement_key_map': self.measurement_key_map,
             'param_resolver': self.param_resolver,
+            'repetition_ids': self.repetition_ids,
         }
 
     @classmethod
     def _from_json_dict_(
-        cls, circuit, repetitions, qubit_map, measurement_key_map, param_resolver, **kwargs
+        cls,
+        circuit,
+        repetitions,
+        qubit_map,
+        measurement_key_map,
+        param_resolver,
+        repetition_ids,
+        **kwargs,
     ):
         return (
             cls(circuit)
             .with_qubit_mapping(dict(qubit_map))
             .with_measurement_key_mapping(measurement_key_map)
             .with_params(param_resolver)
-            .repeat(repetitions)
+            .repeat(repetitions, repetition_ids)
         )
 
     # Methods for constructing a similar object with one field modified.
 
     def repeat(
         self,
-        repetitions: INT_TYPE,
+        repetitions: Optional[INT_TYPE] = None,
+        repetition_ids: Optional[List[str]] = None,
     ) -> 'CircuitOperation':
         """Returns a copy of this operation repeated 'repetitions' times.
+         Each repetition instance will be identified by a single repetition_id.
 
         Args:
             repetitions: Number of times this operation should repeat. This
-                is multiplied with any pre-existing repetitions.
+                is multiplied with any pre-existing repetitions. If unset, it
+                defaults to the length of `repetition_ids`.
+            repetition_ids: List of IDs, one for each repetition. If unset,
+                defaults to `default_repetition_ids(repetitions)`.
 
         Returns:
-            A copy of this operation repeated 'repetitions' times.
+            A copy of this operation repeated `repetitions` times with the
+            appropriate `repetition_ids`. The output `repetition_ids` are the
+            cartesian product of input `repetition_ids` with the base
+            operation's `repetition_ids`. If the base operation has unset
+            `repetition_ids` (indicates {-1, 0, 1} `repetitions` with no custom
+            IDs), the input `repetition_ids` are directly used.
 
         Raises:
             TypeError: `repetitions` is not an integer value.
-            NotImplementedError: The operation contains measurements and
-                cannot have repetitions.
+            ValueError: Unexpected length of `repetition_ids`.
+            ValueError: Both `repetitions` and `repetition_ids` are None.
         """
+        if repetitions is None:
+            if repetition_ids is None:
+                raise ValueError('At least one of repetitions and repetition_ids must be set')
+            repetitions = len(repetition_ids)
+
         if not isinstance(repetitions, (int, np.integer)):
             raise TypeError('Only integer repetitions are allowed.')
-        if repetitions == 1:
+
+        repetitions = int(repetitions)
+
+        if repetitions == 1 and repetition_ids is None:
             # As CircuitOperation is immutable, this can safely return the original.
             return self
-        repetitions = int(repetitions)
-        if protocols.is_measurement(self.circuit):
-            raise NotImplementedError('Loops over measurements are not supported.')
-        return self.replace(repetitions=self.repetitions * repetitions)
+
+        expected_repetition_id_length = abs(repetitions)
+        # The eventual number of repetitions of the returned CircuitOperation.
+        final_repetitions = self.repetitions * repetitions
+
+        if repetition_ids is None:
+            repetition_ids = default_repetition_ids(expected_repetition_id_length)
+        elif len(repetition_ids) != expected_repetition_id_length:
+            raise ValueError(
+                f'Expected repetition_ids={repetition_ids} length to be '
+                f'{expected_repetition_id_length}'
+            )
+
+        # If `self.repetition_ids` is None, this will just return `repetition_ids`.
+        repetition_ids = cartesian_product_of_string_lists(repetition_ids, self.repetition_ids)
+
+        return self.replace(repetitions=final_repetitions, repetition_ids=repetition_ids)
 
     def __pow__(self, power: int) -> 'CircuitOperation':
         return self.repeat(power)
 
+    def with_repetition_ids(self, repetition_ids: List[str]) -> 'CircuitOperation':
+        return self.replace(repetition_ids=repetition_ids)
+
     def with_qubit_mapping(
         self,
         qubit_map: Union[Dict['cirq.Qid', 'cirq.Qid'], Callable[['cirq.Qid'], 'cirq.Qid']],


Test patch:
diff --git a/cirq/circuits/circuit_operation_test.py b/cirq/circuits/circuit_operation_test.py
--- a/cirq/circuits/circuit_operation_test.py
+++ b/cirq/circuits/circuit_operation_test.py
@@ -12,9 +12,12 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+from typing import List
+
 import pytest, sympy
 
 import cirq
+from cirq.circuits.circuit_operation import cartesian_product_of_string_lists
 
 
 def test_properties():
@@ -33,6 +36,7 @@ def test_properties():
     assert op.measurement_key_map == {}
     assert op.param_resolver == cirq.ParamResolver()
     assert op.repetitions == 1
+    assert op.repetition_ids is None
     # Despite having the same decomposition, these objects are not equal.
     assert op != circuit
     assert op == circuit.to_op()
@@ -51,6 +55,32 @@ def test_circuit_type():
         _ = cirq.CircuitOperation(circuit)
 
 
+def test_non_invertible_circuit():
+    a, b, c = cirq.LineQubit.range(3)
+    circuit = cirq.FrozenCircuit(
+        cirq.X(a),
+        cirq.Y(b),
+        cirq.H(c),
+        cirq.CX(a, b) ** sympy.Symbol('exp'),
+        cirq.measure(a, b, c, key='m'),
+    )
+    with pytest.raises(ValueError, match='circuit is not invertible'):
+        _ = cirq.CircuitOperation(circuit, repetitions=-2)
+
+
+def test_repetitions_and_ids_length_mismatch():
+    a, b, c = cirq.LineQubit.range(3)
+    circuit = cirq.FrozenCircuit(
+        cirq.X(a),
+        cirq.Y(b),
+        cirq.H(c),
+        cirq.CX(a, b) ** sympy.Symbol('exp'),
+        cirq.measure(a, b, c, key='m'),
+    )
+    with pytest.raises(ValueError, match='Expected repetition_ids to be a list of length 2'):
+        _ = cirq.CircuitOperation(circuit, repetitions=2, repetition_ids=['a', 'b', 'c'])
+
+
 def test_circuit_sharing():
     a, b, c = cirq.LineQubit.range(3)
     circuit = cirq.FrozenCircuit(
@@ -167,33 +197,68 @@ def test_with_params():
         _ = cirq.resolve_parameters(op_base, cirq.ParamResolver(param_dict))
 
 
-def test_repetition():
+@pytest.mark.parametrize('add_measurements', [True, False])
+@pytest.mark.parametrize('use_default_ids_for_initial_rep', [True, False])
+def test_repeat(add_measurements, use_default_ids_for_initial_rep):
     a, b = cirq.LineQubit.range(2)
-    # This circuit has a modulus of 8.
-    circuit = cirq.FrozenCircuit(cirq.H(a), cirq.CX(a, b))
-    op_base = cirq.CircuitOperation(circuit)
+    circuit = cirq.Circuit(cirq.H(a), cirq.CX(a, b))
+    if add_measurements:
+        circuit.append([cirq.measure(b, key='mb'), cirq.measure(a, key='ma')])
+    op_base = cirq.CircuitOperation(circuit.freeze())
     assert op_base.repeat(1) is op_base
-
-    op_with_reps = op_base.repeat(-5)
-    assert op_with_reps.repetitions == -5
-    assert op_base ** -5 == op_with_reps
-
-    with pytest.raises(TypeError):
+    assert op_base.repeat(1, ['0']) != op_base
+    assert op_base.repeat(1, ['0']) == op_base.repeat(repetition_ids=['0'])
+    assert op_base.repeat(1, ['0']) == op_base.with_repetition_ids(['0'])
+
+    initial_repetitions = -3
+    if add_measurements:
+        with pytest.raises(ValueError, match='circuit is not invertible'):
+            _ = op_base.repeat(initial_repetitions)
+        initial_repetitions = abs(initial_repetitions)
+
+    op_with_reps = None  # type: cirq.CircuitOperation
+    rep_ids = []
+    if use_default_ids_for_initial_rep:
+        op_with_reps = op_base.repeat(initial_repetitions)
+        rep_ids = ['0', '1', '2']
+        assert op_base ** initial_repetitions == op_with_reps
+    else:
+        rep_ids = ['a', 'b', 'c']
+        op_with_reps = op_base.repeat(initial_repetitions, rep_ids)
+        assert op_base ** initial_repetitions != op_with_reps
+        assert (op_base ** initial_repetitions).replace(repetition_ids=rep_ids) == op_with_reps
+    assert op_with_reps.repetitions == initial_repetitions
+    assert op_with_reps.repetition_ids == rep_ids
+    assert op_with_reps.repeat(1) is op_with_reps
+
+    final_repetitions = 2 * initial_repetitions
+
+    op_with_consecutive_reps = op_with_reps.repeat(2)
+    assert op_with_consecutive_reps.repetitions == final_repetitions
+    assert op_with_consecutive_reps.repetition_ids == cartesian_product_of_string_lists(
+        ['0', '1'], rep_ids
+    )
+    assert op_base ** final_repetitions != op_with_consecutive_reps
+
+    op_with_consecutive_reps = op_with_reps.repeat(2, ['a', 'b'])
+    assert op_with_reps.repeat(repetition_ids=['a', 'b']) == op_with_consecutive_reps
+    assert op_with_consecutive_reps.repetitions == final_repetitions
+    assert op_with_consecutive_reps.repetition_ids == cartesian_product_of_string_lists(
+        ['a', 'b'], rep_ids
+    )
+
+    with pytest.raises(ValueError, match='length to be 2'):
+        _ = op_with_reps.repeat(2, ['a', 'b', 'c'])
+
+    with pytest.raises(
+        ValueError, match='At least one of repetitions and repetition_ids must be set'
+    ):
+        _ = op_base.repeat()
+
+    with pytest.raises(TypeError, match='Only integer repetitions are allowed'):
         _ = op_base.repeat(1.3)
 
 
-def test_repeat_measurement_fails():
-    a, b = cirq.LineQubit.range(2)
-    circuit = cirq.FrozenCircuit(
-        cirq.X(a),
-        cirq.measure(b, key='mb'),
-        cirq.measure(a, key='ma'),
-    )
-    op = cirq.CircuitOperation(circuit)
-    with pytest.raises(NotImplementedError):
-        _ = op.repeat(3)
-
-
 def test_qid_shape():
     circuit = cirq.FrozenCircuit(
         cirq.IdentityGate(qid_shape=(q.dimension,)).on(q)
@@ -253,14 +318,16 @@ def test_string_format():
     )
 
     fc2 = cirq.FrozenCircuit(cirq.X(x), cirq.H(y), cirq.CX(y, x))
-    op2 = cirq.CircuitOperation(circuit=fc2, qubit_map=({y: z}), repetitions=3)
+    op2 = cirq.CircuitOperation(
+        circuit=fc2, qubit_map=({y: z}), repetitions=3, repetition_ids=['a', 'b', 'c']
+    )
     assert (
         str(op2)
         == f"""\
 {op2.circuit.serialization_key()}:
 [ 0: ───X───X───          ]
 [           │             ]
-[ 1: ───H───@───          ](qubit_map={{1: 2}}, loops=3)"""
+[ 1: ───H───@───          ](qubit_map={{1: 2}}, repetition_ids=['a', 'b', 'c'])"""
     )
     assert (
         repr(op2)
@@ -277,6 +344,7 @@ def test_string_format():
     ]),
     repetitions=3,
     qubit_map={cirq.LineQubit(1): cirq.LineQubit(2)},
+    repetition_ids=['a', 'b', 'c'],
 )"""
     )
 
@@ -356,6 +424,7 @@ def test_json_dict():
         'qubit_map': sorted([(k, v) for k, v in op.qubit_map.items()]),
         'measurement_key_map': op.measurement_key_map,
         'param_resolver': op.param_resolver,
+        'repetition_ids': None,
     }
 
 
@@ -480,6 +549,30 @@ def test_decompose_loops():
     assert cirq.Circuit(cirq.decompose_once(op)) == expected_circuit
 
 
+def test_decompose_loops_with_measurements():
+    a, b = cirq.LineQubit.range(2)
+    circuit = cirq.FrozenCircuit(
+        cirq.H(a),
+        cirq.CX(a, b),
+        cirq.measure(a, b, key='m'),
+    )
+    base_op = cirq.CircuitOperation(circuit)
+
+    op = base_op.with_qubits(b, a).repeat(3)
+    expected_circuit = cirq.Circuit(
+        cirq.H(b),
+        cirq.CX(b, a),
+        cirq.measure(b, a, key='0-m'),
+        cirq.H(b),
+        cirq.CX(b, a),
+        cirq.measure(b, a, key='1-m'),
+        cirq.H(b),
+        cirq.CX(b, a),
+        cirq.measure(b, a, key='2-m'),
+    )
+    assert cirq.Circuit(cirq.decompose_once(op)) == expected_circuit
+
+
 def test_decompose_nested():
     a, b, c, d = cirq.LineQubit.range(4)
     exp1 = sympy.Symbol('exp1')
@@ -542,6 +635,48 @@ def test_decompose_nested():
     assert cirq.Circuit(cirq.decompose(final_op)) == expected_circuit
 
 
+def test_decompose_repeated_nested_measurements():
+    # Details of this test described at
+    # https://tinyurl.com/measurement-repeated-circuitop#heading=h.sbgxcsyin9wt.
+    a = cirq.LineQubit(0)
+
+    op1 = (
+        cirq.CircuitOperation(cirq.FrozenCircuit(cirq.measure(a, key='A')))
+        .with_measurement_key_mapping({'A': 'B'})
+        .repeat(2, ['zero', 'one'])
+    )
+
+    op2 = (
+        cirq.CircuitOperation(cirq.FrozenCircuit(cirq.measure(a, key='P'), op1))
+        .with_measurement_key_mapping({'B': 'C', 'P': 'Q'})
+        .repeat(2, ['zero', 'one'])
+    )
+
+    op3 = (
+        cirq.CircuitOperation(cirq.FrozenCircuit(cirq.measure(a, key='X'), op2))
+        .with_measurement_key_mapping({'C': 'D', 'X': 'Y'})
+        .repeat(2, ['zero', 'one'])
+    )
+
+    expected_circuit = cirq.Circuit(
+        cirq.measure(a, key='zero-Y'),
+        cirq.measure(a, key='zero-zero-Q'),
+        cirq.measure(a, key='zero-zero-zero-D'),
+        cirq.measure(a, key='zero-zero-one-D'),
+        cirq.measure(a, key='zero-one-Q'),
+        cirq.measure(a, key='zero-one-zero-D'),
+        cirq.measure(a, key='zero-one-one-D'),
+        cirq.measure(a, key='one-Y'),
+        cirq.measure(a, key='one-zero-Q'),
+        cirq.measure(a, key='one-zero-zero-D'),
+        cirq.measure(a, key='one-zero-one-D'),
+        cirq.measure(a, key='one-one-Q'),
+        cirq.measure(a, key='one-one-zero-D'),
+        cirq.measure(a, key='one-one-one-D'),
+    )
+    assert cirq.Circuit(cirq.decompose(op3)) == expected_circuit
+
+
 def test_tag_propagation():
     # Tags are not propagated from the CircuitOperation to its components.
     # TODO: support tag propagation for better serialization.
diff --git a/cirq/protocols/json_test_data/CircuitOperation.json b/cirq/protocols/json_test_data/CircuitOperation.json
--- a/cirq/protocols/json_test_data/CircuitOperation.json
+++ b/cirq/protocols/json_test_data/CircuitOperation.json
@@ -128,7 +128,8 @@
     "param_resolver": {
       "cirq_type": "ParamResolver",
       "param_dict": []
-    }
+    },
+    "repetition_ids": ["0"]
   },
   {
     "cirq_type": "CircuitOperation",
@@ -176,7 +177,8 @@
     "param_resolver": {
       "cirq_type": "ParamResolver",
       "param_dict": []
-    }
+    },
+    "repetition_ids": ["a", "b"]
   },
   {
     "cirq_type": "CircuitOperation",
@@ -224,6 +226,7 @@
           1.5
         ]
       ]
-    }
+    },
+    "repetition_ids": null
   }
 ]
\ No newline at end of file
diff --git a/cirq/protocols/json_test_data/CircuitOperation.repr b/cirq/protocols/json_test_data/CircuitOperation.repr
--- a/cirq/protocols/json_test_data/CircuitOperation.repr
+++ b/cirq/protocols/json_test_data/CircuitOperation.repr
@@ -10,6 +10,7 @@
         cirq.MeasurementGate(5, '0,1,2,3,4', ()).on(cirq.LineQubit(0), cirq.LineQubit(1), cirq.LineQubit(2), cirq.LineQubit(3), cirq.LineQubit(4)),
     ),
 ]),
+repetition_ids=['0'],
 measurement_key_map={'0,1,2,3,4': 'mixed_state'}),
 cirq.CircuitOperation(circuit=cirq.FrozenCircuit([
     cirq.Moment(
@@ -17,6 +18,7 @@ cirq.CircuitOperation(circuit=cirq.FrozenCircuit([
     ),
 ]),
 repetitions=-2,
+repetition_ids=['a', 'b'],
 qubit_map={cirq.LineQubit(0): cirq.LineQubit(1)}),
 cirq.CircuitOperation(circuit=cirq.FrozenCircuit([
     cirq.Moment(