move align step after calculating sweeps as that would lead to mismatched circuit structures in parameterizations

Author: babacry

cirq-core/cirq/transformers/merge_single_qubit_gates.py

@@ -179,7 +179,7 @@ def _sweep_on_symbols(sweep: Sweep, symbols: set[sympy.Symbol]) -> Sweep:
 
 
 def _calc_phxz_sweeps(
-    symbolized_circuit: cirq.Circuit, resolved_circuits: list[cirq.Circuit]
+    symbolized_circuit: cirq.AbstractCircuit, resolved_circuits: list[cirq.AbstractCircuit]
 ) -> Sweep:
     """Return the phxz sweep of the symbolized_circuit on resolved_circuits.
 
@@ -306,18 +306,16 @@ def merge_single_qubit_gates_to_phxz_symbolized(
     ]
 
     # Step 2, get the new symbolized circuit by symbolizing on indexed symbolized_single_tag.
-    new_circuit = align.align_right(
-        tag_transformers.remove_tags(  # remove the temp tags used to track merges
-            symbolize.symbolize_single_qubit_gates_by_indexed_tags(
-                tag_transformers.index_tags(  # index all 1-qubit-ops merged from ops with symbols
-                    merged_circuits[0],
-                    context=transformer_api.TransformerContext(deep=deep),
-                    target_tags={symbolized_single_tag},
-                ),
-                symbolize_tag=symbolize.SymbolizeTag(prefix=symbolized_single_tag),
+    new_circuit = tag_transformers.remove_tags(  # remove the temp tags used to track merges
+        symbolize.symbolize_single_qubit_gates_by_indexed_tags(
+            tag_transformers.index_tags(  # index all 1-qubit-ops merged from ops with symbols
+                merged_circuits[0],
+                context=transformer_api.TransformerContext(deep=deep),
+                target_tags={symbolized_single_tag},
             ),
-            remove_if=lambda tag: str(tag).startswith(symbolized_single_tag),
-        )
+            symbolize_tag=symbolize.SymbolizeTag(prefix=symbolized_single_tag),
+        ),
+        remove_if=lambda tag: str(tag).startswith(symbolized_single_tag),
     )
 
     # Step 3, get N sets of parameterizations as new_sweep.
@@ -326,4 +324,4 @@ def merge_single_qubit_gates_to_phxz_symbolized(
         _sweep_on_symbols(sweep, remaining_symbols),  # remaining sweeps
     )
 
-    return new_circuit, new_sweep
+    return align.align_right(new_circuit), new_sweep

cirq-core/cirq/transformers/merge_single_qubit_gates_test.py

@@ -297,6 +297,29 @@ def test_case1(self):
                 {q: q for q in input_circuit.all_qubits()},
             )
 
+    def test_with_gauge_compiling_as_sweep_success(self):
+        qubits = cirq.LineQubit.range(7)
+        c = cirq.Circuit(
+            cirq.Moment(cirq.H(qubits[0]), cirq.H(qubits[3])),
+            cirq.Moment(cirq.CZ(qubits[0], qubits[2]), cirq.CZ(qubits[3], qubits[5])),
+            cirq.Moment(cirq.CZ(qubits[0], qubits[1]), cirq.CZ(qubits[3], qubits[4])),
+            cirq.Moment(cirq.CZ(qubits[1], qubits[3]), cirq.CZ(qubits[4], qubits[6])),
+            cirq.Moment(cirq.M(*qubits, key='m')),
+        )
+        old_circuit, old_sweep = cirq.transformers.gauge_compiling.CZGaugeTransformer.as_sweep(
+            c, N=50
+        )
+        new_circuit, new_sweep = cirq.merge_single_qubit_gates_to_phxz_symbolized(
+            old_circuit, sweep=old_sweep
+        )
+        # Check the unitaries are preserved for each set of sweep paramerization.
+        for old_resolver, new_resolver in zip(old_sweep, new_sweep):
+            cirq.testing.assert_circuits_have_same_unitary_given_final_permutation(
+                cirq.resolve_parameters(old_circuit[0:-1], old_resolver),
+                cirq.resolve_parameters(new_circuit[0:-1], new_resolver),
+                {q: q for q in qubits},
+            )
+
     def test_case_non_parameterized_singles(self):
         """Test merge_single_qubit_gates_to_phxz_symbolized when all single qubit gates are not
         parameterized."""