Add flag to @cirq.transformer decorator

Author: tanujkhattar

cirq-core/cirq/transformers/transformer_api.py

@@ -20,6 +20,7 @@
 import functools
 import textwrap
 from typing import (
+    cast,
     Any,
     Tuple,
     Hashable,
@@ -32,6 +33,8 @@
 )
 from typing_extensions import Protocol
 
+from cirq import circuits
+
 if TYPE_CHECKING:
     import cirq
 
@@ -214,10 +217,13 @@ class TransformerContext:
                 circuit. Transformers should not transform any operation marked with a tag that
                 belongs to this tuple. Note that any instance of a Hashable type (like `str`,
                 `cirq.VirtualTag` etc.) is a valid tag.
+        deep: If true, the transformer should be recursively applied to all sub-circuits wrapped
+                inside circuit operations.
     """
 
     logger: TransformerLogger = NoOpTransformerLogger()
     tags_to_ignore: Tuple[Hashable, ...] = ()
+    deep: bool = False
 
 
 class TRANSFORMER(Protocol):
@@ -232,16 +238,20 @@ def __call__(
 
 
 @overload
-def transformer(cls_or_func: _TRANSFORMER_T) -> _TRANSFORMER_T:
+def transformer(
+    cls_or_func: _TRANSFORMER_T, *, add_support_for_deep: bool = False
+) -> _TRANSFORMER_T:
     pass
 
 
 @overload
-def transformer(cls_or_func: _TRANSFORMER_CLS_T) -> _TRANSFORMER_CLS_T:
+def transformer(
+    cls_or_func: _TRANSFORMER_CLS_T, *, add_support_for_deep: bool = False
+) -> _TRANSFORMER_CLS_T:
     pass
 
 
-def transformer(cls_or_func: Any) -> Any:
+def transformer(cls_or_func: Any = None, *, add_support_for_deep: bool = False) -> Any:
     """Decorator to verify API and append logging functionality to transformer functions & classes.
 
     A transformer is a callable that takes as inputs a `cirq.AbstractCircuit` and
@@ -284,10 +294,22 @@ def transformer(cls_or_func: Any) -> Any:
 
     Args:
         cls_or_func: The callable class or function to be decorated.
+        add_support_for_deep: If True, the decorator adds the logic to first apply the
+            decorated transformer on subcircuits wrapped inside `cirq.CircuitOperation`s
+            before applying it on the top-level circuit, if context.deep is True.
 
     Returns:
         Decorated class / function which includes additional logging boilerplate.
     """
+
+    # If keyword arguments were specified, python invokes the decorator method
+    # without a `cls` argument, then passes `cls` into the result.
+    if cls_or_func is None:
+        return lambda deferred_cls_or_func: transformer(
+            deferred_cls_or_func,
+            add_support_for_deep=add_support_for_deep,
+        )
+
     if isinstance(cls_or_func, type):
         cls = cls_or_func
         method = cls.__call__
@@ -298,6 +320,7 @@ def method_with_logging(
             self, circuit: 'cirq.AbstractCircuit', **kwargs
         ) -> 'cirq.AbstractCircuit':
             return _transform_and_log(
+                add_support_for_deep,
                 lambda circuit, **kwargs: method(self, circuit, **kwargs),
                 cls.__name__,
                 circuit,
@@ -315,6 +338,7 @@ def method_with_logging(
         @functools.wraps(func)
         def func_with_logging(circuit: 'cirq.AbstractCircuit', **kwargs) -> 'cirq.AbstractCircuit':
             return _transform_and_log(
+                add_support_for_deep,
                 func,
                 func.__name__,
                 circuit,
@@ -325,7 +349,7 @@ def func_with_logging(circuit: 'cirq.AbstractCircuit', **kwargs) -> 'cirq.Abstra
         return func_with_logging
 
 
-def _get_default_context(func: TRANSFORMER):
+def _get_default_context(func: TRANSFORMER) -> TransformerContext:
     sig = inspect.signature(func)
     default_context = sig.parameters["context"].default
     assert (
@@ -334,7 +358,35 @@ def _get_default_context(func: TRANSFORMER):
     return default_context
 
 
+def _run_transformer_on_circuit(
+    add_support_for_deep: bool,
+    func: TRANSFORMER,
+    circuit: 'cirq.AbstractCircuit',
+    extracted_context: Optional[TransformerContext],
+    **kwargs,
+) -> 'cirq.AbstractCircuit':
+    mutable_circuit = None
+    if extracted_context and extracted_context.deep and add_support_for_deep:
+        batch_replace = []
+        for i, op in circuit.findall_operations(
+            lambda o: isinstance(o.untagged, circuits.CircuitOperation)
+        ):
+            op_untagged = cast(circuits.CircuitOperation, op.untagged)
+            if not set(op.tags).isdisjoint(extracted_context.tags_to_ignore):
+                continue
+            op_untagged = op_untagged.replace(
+                circuit=_run_transformer_on_circuit(
+                    add_support_for_deep, func, op_untagged.circuit, extracted_context, **kwargs
+                ).freeze()
+            )
+            batch_replace.append((i, op, op_untagged.with_tags(*op.tags)))
+        mutable_circuit = circuit.unfreeze(copy=True)
+        mutable_circuit.batch_replace(batch_replace)
+    return func(mutable_circuit if mutable_circuit else circuit, **kwargs)
+
+
 def _transform_and_log(
+    add_support_for_deep: bool,
     func: TRANSFORMER,
     transformer_name: str,
     circuit: 'cirq.AbstractCircuit',
@@ -344,7 +396,9 @@ def _transform_and_log(
     """Helper to log initial and final circuits before and after calling the transformer."""
     if extracted_context:
         extracted_context.logger.register_initial(circuit, transformer_name)
-    transformed_circuit = func(circuit, **kwargs)
+    transformed_circuit = _run_transformer_on_circuit(
+        add_support_for_deep, func, circuit, extracted_context, **kwargs
+    )
     if extracted_context:
         extracted_context.logger.register_final(transformed_circuit, transformer_name)
     return transformed_circuit

cirq-core/cirq/transformers/transformer_api_test.py

@@ -21,7 +21,7 @@
 import pytest
 
 
-@cirq.transformer
+@cirq.transformer()
 class MockTransformerClass:
     def __init__(self):
         self.mock = mock.Mock()
@@ -59,6 +59,11 @@ def __call__(
         return circuit[::-1]
 
 
+@cirq.transformer(add_support_for_deep=True)
+class MockTransformerClassSupportsDeep(MockTransformerClass):
+    pass
+
+
 def make_transformer_func_with_defaults() -> cirq.TRANSFORMER:
     my_mock = mock.Mock()
 
@@ -77,16 +82,18 @@ def func(
     return func
 
 
-def make_transformer_func() -> cirq.TRANSFORMER:
+def make_transformer_func(add_support_for_deep: bool = False) -> cirq.TRANSFORMER:
     my_mock = mock.Mock()
 
-    @cirq.transformer
+    @cirq.transformer(add_support_for_deep=add_support_for_deep)
     def mock_tranformer_func(
         circuit: cirq.AbstractCircuit, *, context: Optional[cirq.TransformerContext] = None
     ) -> cirq.Circuit:
+        print("CALLED:", circuit, context, sep="\n\n----------\n\n")
         my_mock(circuit, context)
         return circuit.unfreeze()
 
+    # mock_transformer = mock_tranformer_func(add_support_for_deep)
     mock_tranformer_func.mock = my_mock  # type: ignore
     return mock_tranformer_func
 
@@ -134,6 +141,51 @@ def test_transformer_decorator_with_defaults(transformer):
     transformer.mock.assert_called_with(circuit, context, 1e-2, CustomArg(12))
 
 
+# from unittest.mock import patch, call
+
+
+@pytest.mark.parametrize(
+    'transformer, supports_deep',
+    [
+        (MockTransformerClass(), False),
+        (make_transformer_func(), False),
+        (MockTransformerClassSupportsDeep(), True),
+        (make_transformer_func(add_support_for_deep=True), True),
+    ],
+)
+def test_transformer_decorator_adds_support_for_deep(transformer, supports_deep):
+    q1, q2 = cirq.LineQubit.range(2)
+    c_nested_x = cirq.FrozenCircuit(cirq.X(q1))
+    c_nested_y = cirq.FrozenCircuit(cirq.Y(q1))
+    c_nested_xy = cirq.FrozenCircuit(
+        cirq.CircuitOperation(c_nested_x).repeat(5).with_tags("ignore"),
+        cirq.CircuitOperation(c_nested_y).repeat(7).with_tags("preserve_tag"),
+    )
+    c_nested_yx = cirq.FrozenCircuit(
+        cirq.CircuitOperation(c_nested_y).repeat(7).with_tags("ignore"),
+        cirq.CircuitOperation(c_nested_x).repeat(5).with_tags("preserve_tag"),
+    )
+    c_orig = cirq.Circuit(
+        cirq.CircuitOperation(c_nested_xy).repeat(4),
+        cirq.CircuitOperation(c_nested_x).repeat(5).with_tags("ignore"),
+        cirq.CircuitOperation(c_nested_y).repeat(6),
+        cirq.CircuitOperation(c_nested_yx).repeat(7),
+    )
+    context = cirq.TransformerContext(tags_to_ignore=["ignore"], deep=True)
+    transformer(c_orig, context=context)
+    expected_calls = [mock.call(c_orig, context)]
+    if supports_deep:
+        expected_calls = [
+            mock.call(c_nested_y, context),  # c_orig --> xy --> y
+            mock.call(c_nested_xy, context),  # c_orig --> xy
+            mock.call(c_nested_y, context),  # c_orig --> y
+            mock.call(c_nested_x, context),  # c_orig --> yx --> x
+            mock.call(c_nested_yx, context),  # c_orig --> yx
+            mock.call(c_orig, context),  # c_orig
+        ]
+    transformer.mock.assert_has_calls(expected_calls)
+
+
 @cirq.transformer
 class T1:
     def __call__(