Preserve circuit tags in transformer_primitives.map_operations (quantumlib#6505)

Review: @NoureldinYosri

Author: maffoo

cirq/transformers/transformer_primitives.py

@@ -53,12 +53,9 @@ def _to_target_circuit_type(
 
 
 def _create_target_circuit_type(ops: ops.OP_TREE, target_circuit: CIRCUIT_TYPE) -> CIRCUIT_TYPE:
-    return cast(
-        CIRCUIT_TYPE,
-        circuits.Circuit(ops)
-        if isinstance(target_circuit, circuits.Circuit)
-        else circuits.FrozenCircuit(ops),
-    )
+    if isinstance(target_circuit, circuits.FrozenCircuit):
+        return cast(CIRCUIT_TYPE, circuits.FrozenCircuit(ops).with_tags(*target_circuit.tags))
+    return cast(CIRCUIT_TYPE, circuits.Circuit(ops))
 
 
 def map_moments(

cirq/transformers/transformer_primitives_test.py

@@ -205,6 +205,33 @@ def map_func(op: cirq.Operation, _: int) -> cirq.OP_TREE:
 # pylint: enable=line-too-long
 
 
+@pytest.mark.parametrize("deep", [False, True])
+def test_map_operations_preserves_circuit_tags(deep: bool) -> None:
+    tag = "should be preserved"
+
+    def func(op: cirq.Operation, idx: int) -> cirq.Operation:
+        return cirq.Y(op.qubits[0]) if op.gate == cirq.X else op
+
+    x = cirq.X(cirq.q(0))
+    circuit = cirq.FrozenCircuit.from_moments(x, cirq.FrozenCircuit(x)).with_tags(tag)
+    mapped = cirq.map_operations(circuit, func, deep=deep)
+
+    assert mapped.tags == (tag,)
+
+
+def test_map_operations_deep_preserves_subcircuit_tags():
+    tag = "should be preserved"
+
+    def func(op: cirq.Operation, idx: int) -> cirq.Operation:
+        return cirq.Y(op.qubits[0]) if op.gate == cirq.X else op
+
+    x = cirq.X(cirq.q(0))
+    circuit = cirq.FrozenCircuit.from_moments(x, cirq.FrozenCircuit(x).with_tags(tag))
+    mapped = cirq.map_operations(circuit, func, deep=True)
+
+    assert mapped[1].operations[0].circuit.tags == (tag,)
+
+
 def test_map_operations_deep_respects_tags_to_ignore():
     q = cirq.LineQubit.range(2)
     c_nested = cirq.FrozenCircuit(cirq.CX(*q), cirq.CX(*q).with_tags("ignore"), cirq.CX(*q))