Fix behavior of MergeInteractions for tagged partial CZ gates (quantumlib#4288)

* Add test for the failing case with a tagged operation

* Fix merge interactions for tagged operation

Author: cduck

cirq/optimizers/merge_interactions.py

@@ -55,16 +55,9 @@ def optimization_at(
 
         switch_to_new = False
         switch_to_new |= any(
-            len(old_op.qubits) == 2 and not isinstance(old_op.gate, ops.CZPowGate)
+            len(old_op.qubits) == 2 and not self._may_keep_old_op(old_op)
             for old_op in old_operations
         )
-        if not self.allow_partial_czs:
-            switch_to_new |= any(
-                isinstance(old_op, ops.GateOperation)
-                and isinstance(old_op.gate, ops.CZPowGate)
-                and old_op.gate.exponent != 1
-                for old_op in old_operations
-            )
 
         # This point cannot be optimized using this method
         if not switch_to_new and old_interaction_count <= 1:
@@ -89,6 +82,13 @@ def optimization_at(
             new_operations=new_operations,
         )
 
+    def _may_keep_old_op(self, old_op: 'cirq.Operation') -> bool:
+        """Returns True if the old two-qubit operation may be left unchanged
+        without decomposition."""
+        if self.allow_partial_czs:
+            return isinstance(old_op.gate, ops.CZPowGate)
+        return isinstance(old_op.gate, ops.CZPowGate) and old_op.gate.exponent == 1
+
     def _op_to_matrix(
         self, op: ops.Operation, qubits: Tuple['cirq.Qid', ...]
     ) -> Optional[np.ndarray]:

cirq/optimizers/merge_interactions_test.py

@@ -163,6 +163,16 @@ def test_optimizes_single_iswap():
     assert len([1 for op in c.all_operations() if len(op.qubits) == 2]) == 2
 
 
+def test_optimizes_tagged_partial_cz():
+    a, b = cirq.LineQubit.range(2)
+    c = cirq.Circuit((cirq.CZ ** 0.5)(a, b).with_tags('mytag'))
+    assert_optimization_not_broken(c)
+    cirq.MergeInteractions(allow_partial_czs=False).optimize_circuit(c)
+    assert (
+        len([1 for op in c.all_operations() if len(op.qubits) == 2]) == 2
+    ), 'It should take 2 CZ gates to decompose a CZ**0.5 gate'
+
+
 @pytest.mark.parametrize(
     'circuit',
     (