Implement a way to identify CircuitOperations (operations with loops)

[dkafri]

Is your feature request related to a use case or problem? Please describe.
I've been doing a lot of transformations on Circuits containing CircuitOperations. Sometimes I need to manipulate CircuitOperations. The way I nominally identify them is using isinstance(..., cirq.CircuitOperation). But it turns out this doesn't work if the CircuitOperation is contained in a cirq.TaggedOperation.

Describe the solution you'd like
Implement a function like cirq.measurement_key_objs that I can trust to always tell me if an operation is a CircuitOperation (or effectively one under the hood).

[optional] Describe alternatives/workarounds you've considered
Writing my own custom identifier that handles TaggedOperations. But am I missing something?

[optional] Additional context (e.g. screenshots)

What is the urgency from your perspective for this issue? Is it blocking important work?

P3 - I'm not really blocked by it, it is an idea I'd like to discuss / suggestion based on principle

[maffoo]

The incantation I usually use is isinstance(op.untagged, cirq.CircuitOperation). But we could certainly add a helper function for this.

[dkafri]

It becomes somewhat cumbersome, since I usually want to use CircuitOperation.replace, and then I have to worry about keeping track of whether I need to use the untagged op or not.

[maffoo]

I agree this can be a bit tedious. Often typing will help with this, since if you do isinstance(op.untagged, cirq.CircuitOperation) then the type checker will know that op.untagged has a replace method, whereas op, which is of type cirq.Operation will not have a replace method. You do have to remember to call with_tags to add back the tags after manipulating the circuit, which is easy to forget. I guess it would help to clarify the full extent of what you're trying to do since the issue title only asks how to "identify" circuit operations, but it sounds like you also want to transform them. We could certainly think about some helpers to make both of these easier.

Note that the root problem here is that tags are handled by wrapping operations with TaggedOperation which changes the type, so this issue affects much more than just manipulations on CircuitOperation, it can affect a lot of code that wants to find and manipulate operations. (For gates it is often simpler since TaggedOperation has a pass-through .gate property to get the underlying gate, but we don't have something like this for circuit operations.) We took a different approach with FrozenCircuit tags, where tag support is built into the class itself rather than using a wrapper. It might be possible to change Operation to include tagging support directly, we would just have to be careful about the various subclasses, though I think there are not too many subclasses.

[mathe-matician]

@pavoljuhas @verult I'm interested in helping out on this issue if it is still available.

@maffoo your comment provides some helpful direction for this issue - is the ask here still just to add a helper function, or is the scope a bit larger to add some tag support to the base cirq.Operation class? The latter seems like a better solution long term to me as well.

[daxfohl]

xref #3678 #4193

[daxfohl]

tbh I'd consider removing immutability from the next version of Cirq. Immutability is nice for lots of things, but I don't feel like it's providing all that much value here. It affects circuit construction perf, and makes things like this unnecessarily painful. It keeps people from shooting themselves in the foot in some ways, but makes things harder to work with in other ways. If we remove immutability, we can just have Operation.set_tags and be done. Any design to do so while retaining immutability and open subclasses seems to lead to suboptimal tradeoffs; either what we have already, or some use of copy dunder methods, or having to implement the same function independently in every subclass.

[pavoljuhas]

@mathe-matician - thanks for volunteering to help!

@dkafri - I wonder if it would work to add some replace_untagged(self, transform: Callable[Operation, Operation]) convenience helper to the Operation class which would transform the operation and reapply tags that were there initially?

[dkafri]

@mathe-matician - thanks for volunteering to help!

@dkafri - I wonder if it would work to add some replace_untagged(self, transform: Callable[Operation, Operation]) convenience helper to the Operation class which would transform the operation and reapply tags that were there initially?

Do you mean to add this to a Circuit or CircuitOperation? It feels odd to add this directly to Operation (but maybe I am just confused). Here is what I'm currently doing at the level of Circuits (to deal with CircuitOperations, not tags):

@cirq.transformer(add_deep_support=True)
@attrs.define
class _ReplaceOpWith:
    """cirq Transformer protocol for replacing operations."""

    predicate: Callable[[cirq.Operation], bool]
    replacement: Callable[[cirq.Operation], cirq.OP_TREE]

    def __call__(
        self, circuit: cirq.AbstractCircuit, *, context: cirq.TransformerContext | None = None
    ) -> cirq.Circuit:
        circuit = circuit.unfreeze(copy=True)
        return circuit.map_operations(self._replace_if_pred)

    def _replace_if_pred(self, op: cirq.Operation) -> cirq.Operation:
        return self.replacement(op) if self.predicate(op) else op


def deep_replace(
    circuit: cirq.Circuit | cirq.FrozenCircuit,
    predicate: Callable[[cirq.Operation], bool],
    replacement: Callable[[cirq.Operation], cirq.OP_TREE],
) -> cirq.Circuit:
    """Replace operations in a circuit according to a standard rule.

    This implementation works for nested CircuitOperations.

    Args:
        circuit: The circuit to modify.
        predicate: Operations flagged by this boolean function are modified.
        replacement: Takes an operation to be modified and returns the desired changed operation(s).

    """

    replacer = _ReplaceOpWith(predicate, replacement)
    return replacer(circuit, context=cirq.TransformerContext(deep=True))

[dkafri]

tbh I'd consider removing immutability from the next version of Cirq. Immutability is nice for lots of things, but I don't feel like it's providing all that much value here. It affects circuit construction perf, and makes things like this unnecessarily painful. It keeps people from shooting themselves in the foot in some ways, but makes things harder to work with in other ways. If we remove immutability, we can just have Operation.set_tags and be done. Any design to do so while retaining immutability and open subclasses seems to lead to suboptimal tradeoffs; either what we have already, or some use of copy dunder methods, or having to implement the same function independently in every subclass.

I would be wary of removing immutability from Operations. For example, this would make hashing of Operations impossible (and that is something I depend on currently).

[daxfohl]

Is the transformer library not already handling these? It looks from the code like it should be working around tags on CircuitOperations here.

Cirq/cirq-core/cirq/transformers/transformer_primitives.py

Line 165 in 614c78a

op = op.untagged.replace(