Fix moment commutation detection for group-commuting operations (#6659) (#7082)

* Fix Moment.commutes_ to handle multi-qubit unitaries (#6659)

Add pairwise commuting checks and, when inconclusive, fall back to full
                  unitary commutator. Includes new tests showing Z⊗Z commutes with RXX.
                  Fixes #6659.

* Put back pre-existing commutation tests

* prune duplicate or redundant test cases
* reuse cirq.AmplitudeDampingChannel to test with non-unitary
* move commute-related test functions together

* Refactor `Moment._commutes_` to convert to CircuitOperation if necessary

Also check `measurement_key_objs` and `control_keys` active for checked
Moments as in the Operation commutes protocol.

* Test commutation of moments with measurement keys and controls

Adapt `cirq.ops.classically_controlled_operation_test.test_commute`

* Small tweak of docstring

* Add shared internal function _operations_commutes_impl

Generalize `Operation._commutes_` in a shared internal function
`_operations_commutes_impl`

* Sync module variable name with circuit module name

No change in code function.

* Use shared _operations_commutes_impl for Moment._commutes_

Also adjust `_operations_commutes_impl` to assume equal Moment-s
and Operation-s commute.

* Few optimizations in _operations_commutes_impl

- frozenset().union(*tuples) is faster than iterating over all
  items in every tuple to construct a frozenset.

- base class Operation does not have `__hash__` method, therefore
  we cannot use `set(operations)`.

---------

Co-authored-by: Pavol Juhas <juhas@google.com>

Author: chamodtharakaperera

cirq-core/cirq/circuits/moment.py

@@ -49,7 +49,7 @@
     import cirq
 
 # Lazy imports to break circular dependencies.
-circuits = LazyLoader("circuits", globals(), "cirq.circuits.circuit")
+circuit = LazyLoader("circuit", globals(), "cirq.circuits.circuit")
 op_tree = LazyLoader("op_tree", globals(), "cirq.ops.op_tree")
 text_diagram_drawer = LazyLoader(
     "text_diagram_drawer", globals(), "cirq.circuits.text_diagram_drawer"
@@ -525,7 +525,7 @@ def _from_json_dict_(cls, operations, **kwargs):
         return cls.from_ops(*operations)
 
     def __add__(self, other: 'cirq.OP_TREE') -> 'cirq.Moment':
-        if isinstance(other, circuits.AbstractCircuit):
+        if isinstance(other, circuit.AbstractCircuit):
             return NotImplemented  # Delegate to Circuit.__radd__.
         return self.with_operations(other)
 
@@ -659,37 +659,26 @@ def cleanup_key(key: Any) -> Any:
         return diagram.render()
 
     def _commutes_(self, other: Any, *, atol: float = 1e-8) -> Union[bool, NotImplementedType]:
-        """Determines whether Moment commutes with the Operation.
+        """Determines whether Moment commutes with the other Moment or Operation.
 
         Args:
-            other: An Operation object. Other types are not implemented yet.
-                In case a different type is specified, NotImplemented is
-                returned.
+            other: An Operation or Moment object to test for commutativity.
             atol: Absolute error tolerance. If all entries in v1@v2 - v2@v1
                 have a magnitude less than this tolerance, v1 and v2 can be
                 reported as commuting. Defaults to 1e-8.
 
         Returns:
-            True: The Moment and Operation commute OR they don't have shared
-            quibits.
+            True: The Moment commutes with Moment or Operation OR they don't
+                have shared qubits.
             False: The two values do not commute.
             NotImplemented: In case we don't know how to check this, e.g.
-                the parameter type is not supported yet.
+                the parameter type is not supported or commutativity cannot be
+                determined.
         """
-        if not isinstance(other, ops.Operation):
+        if not isinstance(other, (ops.Operation, Moment)):
             return NotImplemented
-
-        other_qubits = set(other.qubits)
-        for op in self.operations:
-            if not other_qubits.intersection(set(op.qubits)):
-                continue
-
-            commutes = protocols.commutes(op, other, atol=atol, default=NotImplemented)
-
-            if not commutes or commutes is NotImplemented:
-                return commutes
-
-        return True
+        other_operations = other.operations if isinstance(other, Moment) else (other,)
+        return raw_types._operations_commutes_impl(self.operations, other_operations, atol=atol)
 
 
 class _SortByValFallbackToType:

cirq-core/cirq/circuits/moment_test.py

@@ -680,15 +680,15 @@ def test_text_diagram_does_not_depend_on_insertion_order():
     assert str(m1) == str(m2)
 
 
-def test_commutes():
+def test_commutes_moment_and_operation():
     a = cirq.NamedQubit('a')
     b = cirq.NamedQubit('b')
     c = cirq.NamedQubit('c')
     d = cirq.NamedQubit('d')
 
     moment = cirq.Moment([cirq.X(a), cirq.Y(b), cirq.H(c)])
 
-    assert NotImplemented == cirq.commutes(moment, a, default=NotImplemented)
+    assert cirq.commutes(moment, a, default=None) is None
 
     assert cirq.commutes(moment, cirq.X(a))
     assert cirq.commutes(moment, cirq.Y(b))
@@ -700,6 +700,101 @@ def test_commutes():
     assert not cirq.commutes(moment, cirq.H(b))
     assert not cirq.commutes(moment, cirq.X(c))
 
+    # Empty moment commutes with everything
+    moment = cirq.Moment()
+    assert cirq.commutes(moment, cirq.X(a))
+    assert cirq.commutes(moment, cirq.measure(b))
+
+    # Two qubit operation
+    moment = cirq.Moment(cirq.Z(a), cirq.Z(b))
+    assert cirq.commutes(moment, cirq.XX(a, b))
+
+
+def test_commutes_moment_and_moment():
+    a = cirq.NamedQubit('a')
+    b = cirq.NamedQubit('b')
+    c = cirq.NamedQubit('c')
+
+    # Test cases where individual operations don't commute but moments do
+    # Two Z gates (Z⊗Z) commutes with RXX even though individual Z's don't
+    assert not cirq.commutes(cirq.Moment(cirq.Z(a)), cirq.Moment(cirq.XX(a, b)))
+    assert cirq.commutes(cirq.Moment(cirq.Z(a), cirq.Z(b)), cirq.Moment(cirq.XX(a, b)))
+
+    # Moments that do not commute if acting on same qubits
+    assert cirq.commutes(cirq.Moment(cirq.X(a)), cirq.Moment(cirq.Y(b)))
+    assert not cirq.commutes(cirq.Moment(cirq.X(a)), cirq.Moment(cirq.Y(a)))
+
+    # Moments commute with themselves
+    assert cirq.commutes(
+        cirq.Moment([cirq.X(a), cirq.Y(b), cirq.H(c)]),
+        cirq.Moment([cirq.X(a), cirq.Y(b), cirq.H(c)]),
+    )
+
+
+def test_commutes_moment_with_controls():
+    a, b = cirq.LineQubit.range(2)
+    assert cirq.commutes(
+        cirq.Moment(cirq.measure(a, key='k0')), cirq.Moment(cirq.X(b).with_classical_controls('k1'))
+    )
+    assert cirq.commutes(
+        cirq.Moment(cirq.X(b).with_classical_controls('k1')), cirq.Moment(cirq.measure(a, key='k0'))
+    )
+    assert cirq.commutes(
+        cirq.Moment(cirq.X(a).with_classical_controls('k0')),
+        cirq.Moment(cirq.H(b).with_classical_controls('k0')),
+    )
+    assert cirq.commutes(
+        cirq.Moment(cirq.X(a).with_classical_controls('k0')),
+        cirq.Moment(cirq.X(a).with_classical_controls('k0')),
+    )
+    assert not cirq.commutes(
+        cirq.Moment(cirq.measure(a, key='k0')), cirq.Moment(cirq.X(b).with_classical_controls('k0'))
+    )
+    assert not cirq.commutes(
+        cirq.Moment(cirq.X(b).with_classical_controls('k0')), cirq.Moment(cirq.measure(a, key='k0'))
+    )
+    assert not cirq.commutes(
+        cirq.Moment(cirq.X(a).with_classical_controls('k0')),
+        cirq.Moment(cirq.H(a).with_classical_controls('k0')),
+    )
+
+
+def test_commutes_moment_and_moment_comprehensive():
+    a, b, c, d = cirq.LineQubit.range(4)
+
+    # Basic Z⊗Z commuting with XX at different angles
+    m1 = cirq.Moment([cirq.Z(a), cirq.Z(b)])
+    m2 = cirq.Moment([cirq.XXPowGate(exponent=0.5)(a, b)])
+    assert cirq.commutes(m1, m2)
+
+    # Disjoint qubit sets
+    m1 = cirq.Moment([cirq.X(a), cirq.Y(b)])
+    m2 = cirq.Moment([cirq.Z(c), cirq.H(d)])
+    assert cirq.commutes(m1, m2)
+
+    # Mixed case - some commute individually, some as group
+    m1 = cirq.Moment([cirq.Z(a), cirq.Z(b), cirq.X(c)])
+    m2 = cirq.Moment([cirq.XXPowGate(exponent=0.5)(a, b), cirq.X(c)])
+    assert cirq.commutes(m1, m2)
+
+    # Non-commuting case: X on first qubit, Z on second with XX gate
+    m1 = cirq.Moment([cirq.X(a), cirq.Z(b)])
+    m2 = cirq.Moment([cirq.XX(a, b)])
+    assert not cirq.commutes(m1, m2)
+
+    # Complex case requiring unitary calculation - non-commuting case
+    m1 = cirq.Moment([cirq.Z(a), cirq.Z(b), cirq.Z(c)])
+    m2 = cirq.Moment([cirq.XXPowGate(exponent=0.5)(a, b), cirq.X(c)])
+    assert not cirq.commutes(m1, m2)  # Z⊗Z⊗Z doesn't commute with XX⊗X
+
+
+def test_commutes_handles_non_unitary_operation():
+    a = cirq.NamedQubit('a')
+    op_damp_a = cirq.AmplitudeDampingChannel(gamma=0.1).on(a)
+    assert cirq.commutes(cirq.Moment(cirq.X(a)), op_damp_a, default=None) is None
+    assert cirq.commutes(cirq.Moment(cirq.X(a)), cirq.Moment(op_damp_a), default=None) is None
+    assert cirq.commutes(cirq.Moment(op_damp_a), cirq.Moment(op_damp_a))
+
 
 def test_transform_qubits():
     a, b = cirq.LineQubit.range(2)

cirq-core/cirq/ops/raw_types.py

@@ -688,41 +688,7 @@ def _commutes_(
         """Determine if this Operation commutes with the object"""
         if not isinstance(other, Operation):
             return NotImplemented
-
-        self_keys = protocols.measurement_key_objs(self)
-        other_keys = protocols.measurement_key_objs(other)
-        if (
-            not self_keys.isdisjoint(other_keys)
-            or not protocols.control_keys(self).isdisjoint(other_keys)
-            or not protocols.control_keys(other).isdisjoint(self_keys)
-        ):
-            return False
-
-        if hasattr(other, 'qubits') and set(self.qubits).isdisjoint(other.qubits):
-            return True
-
-        from cirq import circuits
-
-        # Remove the classical controls to validate the quantum commutativity. This can be done
-        # because during execution, the two operations will either both be run, in which case they
-        # behave like the suboperations, so if the suboperations commute then these commute. Or
-        # one of them is cold in which case it behaves like the identity, which always commutes.
-        self_raw = self.without_classical_controls()
-        other_raw = other.without_classical_controls()
-        circuit12 = circuits.Circuit(self_raw, other_raw)
-        circuit21 = circuits.Circuit(other_raw, self_raw)
-
-        # Don't create gigantic matrices.
-        shape = protocols.qid_shape_protocol.qid_shape(circuit12)
-        if np.prod(shape, dtype=np.int64) > 2**10:
-            return NotImplemented  # pragma: no cover
-
-        m12 = protocols.unitary_protocol.unitary(circuit12, default=None)
-        m21 = protocols.unitary_protocol.unitary(circuit21, default=None)
-        if m12 is None or m21 is None:
-            return NotImplemented
-
-        return np.allclose(m12, m21, atol=atol)
+        return _operations_commutes_impl([self], [other], atol=atol)
 
     @property
     def classical_controls(self) -> FrozenSet['cirq.Condition']:
@@ -1112,3 +1078,78 @@ def _validate_qid_shape(val: Any, qubits: Sequence['cirq.Qid']) -> None:
         raise ValueError(
             f'Duplicate qids for <{val!r}>. Expected unique qids but got <{qubits!r}>.'
         )
+
+
+def _operations_commutes_impl(
+    ops1: Collection[Operation], ops2: Collection[Operation], *, atol: float
+) -> Union[bool, NotImplementedType]:
+    """Determine if two collections of non-overlapping Operations commute.
+
+    This function implements the commutes protocol for the Operation and Moment classes
+    and is not intended for other use.
+
+    Args:
+        ops1: The first collection of operations.  It is assumed each operation
+            acts on different qubits, i.e., the operations can form a Moment.
+        ops2: The second collection of operations to be checked for commutation
+            with `ops1`.  It is assumed each operation acts on different qubits,
+            i.e., the operations can form a Moment.
+        atol: Absolute error tolerance. If all entries in ops1@ops2 - ops2@ops1
+            have a magnitude less than this tolerance, ops1 and ops2 are considered
+            to commute.
+
+    Returns:
+        True: `ops1` and `ops2` commute (or approximately commute).
+        False: `ops1` and `ops2` do not commute.
+        NotImplemented: The commutativity cannot be determined here.
+    """
+    ops1_keys = frozenset(k for op in ops1 for k in protocols.measurement_key_objs(op))
+    ops2_keys = frozenset(k for op in ops2 for k in protocols.measurement_key_objs(op))
+    ops1_control_keys = frozenset(k for op in ops1 for k in protocols.control_keys(op))
+    ops2_control_keys = frozenset(k for op in ops2 for k in protocols.control_keys(op))
+    if (
+        not ops1_keys.isdisjoint(ops2_keys)
+        or not ops1_control_keys.isdisjoint(ops2_keys)
+        or not ops2_control_keys.isdisjoint(ops1_keys)
+    ):
+        return False
+
+    ops1_qubits = frozenset().union(*(op.qubits for op in ops1))
+    ops2_qubits = frozenset().union(*(op.qubits for op in ops2))
+    if ops1_qubits.isdisjoint(ops2_qubits):
+        return True
+
+    from cirq import circuits
+
+    # Remove the classical controls to validate the quantum commutativity. This can be done
+    # because during execution, the two operations will either both be run, in which case they
+    # behave like the suboperations, so if the suboperations commute then these commute. Or
+    # one of them is cold in which case it behaves like the identity, which always commutes.
+    shared_qubits = ops1_qubits.intersection(ops2_qubits)
+    ops1_raw = [
+        op.without_classical_controls() for op in ops1 if not shared_qubits.isdisjoint(op.qubits)
+    ]
+    ops2_raw = [
+        op.without_classical_controls() for op in ops2 if not shared_qubits.isdisjoint(op.qubits)
+    ]
+    moment1 = circuits.Moment(ops1_raw)
+    moment2 = circuits.Moment(ops2_raw)
+
+    # shortcut if we have equal moments
+    if moment1 == moment2:
+        return True
+
+    circuit12 = circuits.Circuit(moment1, moment2)
+    circuit21 = circuits.Circuit(moment2, moment1)
+
+    # Don't create gigantic matrices.
+    shape = protocols.qid_shape_protocol.qid_shape(circuit12)
+    if np.prod(shape, dtype=np.int64) > 2**10:
+        return NotImplemented  # pragma: no cover
+
+    m12 = protocols.unitary_protocol.unitary(circuit12, default=None)
+    m21 = protocols.unitary_protocol.unitary(circuit21, default=None)
+    if m12 is None or m21 is None:
+        return NotImplemented
+
+    return np.allclose(m12, m21, atol=atol)