Add support for deep=True to cirq.stratified_circuit transformer (#…

…5117)

- Adds support to recursively run `cirq.stratified_circuit` transformer on circuits wrapped inside a circuit operation by setting deep=True in transformer context.
- Part of #5039

cirq-core/cirq/transformers/stratify.py

@@ -45,7 +45,7 @@
 ]
 
 
-@transformer_api.transformer
+@transformer_api.transformer(add_deep_support=True)
 def stratified_circuit(
     circuit: 'cirq.AbstractCircuit',
     *,

cirq-core/cirq/transformers/stratify_test.py

@@ -295,6 +295,43 @@ def test_complex_circuit():
     )
 
 
+def test_complex_circuit_deep():
+    q = cirq.LineQubit.range(5)
+    c_nested = cirq.FrozenCircuit(
+        cirq.Moment(
+            cirq.X(q[0]).with_tags("ignore"),
+            cirq.ISWAP(q[1], q[2]).with_tags("ignore"),
+            cirq.Z(q[4]),
+        ),
+        cirq.Moment(cirq.Z(q[1]), cirq.ISWAP(q[3], q[4])),
+        cirq.Moment(cirq.ISWAP(q[0], q[1]), cirq.X(q[3])),
+        cirq.Moment(cirq.X.on_each(q[0])),
+    )
+    c_nested_stratified = cirq.FrozenCircuit(
+        cirq.Moment(cirq.X(q[0]).with_tags("ignore"), cirq.ISWAP(q[1], q[2]).with_tags("ignore")),
+        cirq.Moment(cirq.Z.on_each(q[1], q[4])),
+        cirq.Moment(cirq.ISWAP(*q[:2]), cirq.ISWAP(*q[3:])),
+        cirq.Moment(cirq.X.on_each(q[0], q[3])),
+    )
+    c_orig = cirq.Circuit(
+        c_nested,
+        cirq.CircuitOperation(c_nested).repeat(5).with_tags("ignore"),
+        c_nested,
+        cirq.CircuitOperation(c_nested).repeat(6).with_tags("preserve_tag"),
+        c_nested,
+    )
+    c_expected = cirq.Circuit(
+        c_nested_stratified,
+        cirq.CircuitOperation(c_nested).repeat(5).with_tags("ignore"),
+        c_nested_stratified,
+        cirq.CircuitOperation(c_nested_stratified).repeat(6).with_tags("preserve_tag"),
+        c_nested_stratified,
+    )
+    context = cirq.TransformerContext(tags_to_ignore=["ignore"], deep=True)
+    c_stratified = cirq.stratified_circuit(c_orig, context=context, categories=[cirq.X, cirq.Z])
+    cirq.testing.assert_same_circuits(c_stratified, c_expected)
+
+
 def test_no_categories_earliest_insert():
     q1, q2, q3, q4, q5 = cirq.LineQubit.range(5)
     input_circuit = cirq.Circuit(