Add handling for sympy conditions in deferred measurement transformer

cirq-core/cirq/transformers/measurement_transformers.py

@@ -13,7 +13,7 @@
 # limitations under the License.
 
 import itertools
-from typing import Any, Dict, List, Optional, TYPE_CHECKING, Union
+from typing import Any, Dict, Iterable, List, Mapping, Optional, TYPE_CHECKING, Union
 
 from cirq import ops, protocols, value
 from cirq.transformers import transformer_api, transformer_primitives
@@ -82,7 +82,6 @@ def defer_measurements(
         A circuit with equivalent logic, but all measurements at the end of the
         circuit.
     Raises:
-        ValueError: If sympy-based classical conditions are used.
         NotImplementedError: When attempting to defer a measurement with a
             confusion map. (https://github.com/quantumlib/Cirq/issues/5482)
     """
@@ -112,19 +111,28 @@ def defer(op: 'cirq.Operation', _) -> 'cirq.OP_TREE':
         elif op.classical_controls:
             new_op = op.without_classical_controls()
             for c in op.classical_controls:
-                if isinstance(c, value.KeyCondition):
-                    if c.key not in measurement_qubits:
-                        raise ValueError(f'Deferred measurement for key={c.key} not found.')
-                    qs = measurement_qubits[c.key]
-                    if len(qs) == 1:
-                        control_values: Any = range(1, qs[0].dimension)
-                    else:
-                        all_values = itertools.product(*[range(q.dimension) for q in qs])
-                        anything_but_all_zeros = tuple(itertools.islice(all_values, 1, None))
-                        control_values = ops.SumOfProducts(anything_but_all_zeros)
-                    new_op = new_op.controlled_by(*qs, control_values=control_values)
+                # Convert to a quantum control
+                missing_keys = [k for k in c.keys if k not in measurement_qubits]
+                if missing_keys:
+                    raise ValueError(f'Deferred measurement for key={missing_keys[0]} not found.')
+                qs = [q for k in c.keys for q in measurement_qubits[k]]
+
+                # Try every possible datastore state against the condition, and the ones that work
+                # are the control values for the new op.
+                datastores = _all_possible_datastore_states(c.keys, measurement_qubits)
+                compatible_datastores = [store for store in datastores if c.resolve(store)]
+
+                # Rearrange these into the format expected by SumOfProducts
+                products = [
+                    [i for k in c.keys for i in store.records[k][0]]
+                    for store in compatible_datastores
+                ]
+                if len(qs) == 1:
+                    # Convenience: this renders more nicely than SumOfProducts.
+                    control_values: Any = [[x[0] for x in products]]
                 else:
-                    raise ValueError('Only KeyConditions are allowed.')
+                    control_values = ops.SumOfProducts(products)
+                new_op = new_op.controlled_by(*qs, control_values=control_values)
             return new_op
         return op
 
@@ -139,6 +147,37 @@ def defer(op: 'cirq.Operation', _) -> 'cirq.OP_TREE':
     return circuit
 
 
+def _all_possible_datastore_states(
+    keys: Iterable['cirq.MeasurementKey'],
+    measurement_qubits: Mapping['cirq.MeasurementKey', Iterable['cirq.Qid']],
+) -> Iterable['cirq.ClassicalDataStoreReader']:
+    """The full product of possible DatsStores states for the given keys."""
+    # First we get the list of all possible values. So if we have a key mapped to qubits of shape
+    # (2, 2) and a key mapped to a qutrit, the possible measurement values are:
+    # [((0, 0), (0,)),
+    #  ((0, 0), (1,)),
+    #  ((0, 0), (2,)),
+    #  ((0, 1), (0,)),
+    #  ((0, 1), (1,)),
+    #  ((0, 1), (2,)),
+    #  ((1, 0), (0,)),
+    #  ((1, 0), (1,)),
+    #  ((1, 0), (2,)),
+    #  ((1, 1), (0,)),
+    #  ((1, 1), (1,)),
+    #  ((1, 1), (2,))]
+    all_values = itertools.product(
+        *[
+            tuple(itertools.product(*[range(q.dimension) for q in measurement_qubits[k]]))
+            for k in keys
+        ]
+    )
+    # Then we create the ClassicalDataDictionaryStore for each of the above.
+    for sequences in all_values:
+        lookup = {k: [sequence] for k, sequence in zip(keys, sequences)}
+        yield value.ClassicalDataDictionaryStore(_records=lookup)
+
+
 @transformer_api.transformer
 def dephase_measurements(
     circuit: 'cirq.AbstractCircuit',

cirq-core/cirq/transformers/measurement_transformers_test.py

@@ -15,6 +15,7 @@
 import numpy as np
 import pytest
 import sympy
+from sympy.parsing import sympy_parser
 
 import cirq
 from cirq.transformers.measurement_transformers import _MeasurementQid
@@ -58,6 +59,79 @@ def test_basic():
     )
 
 
+def test_sympy_control():
+    q0, q1 = cirq.LineQubit.range(2)
+    circuit = cirq.Circuit(
+        cirq.measure(q0, key='a'),
+        cirq.X(q1).with_classical_controls(sympy.Symbol('a')),
+        cirq.measure(q1, key='b'),
+    )
+    assert_equivalent_to_deferred(circuit)
+    deferred = cirq.defer_measurements(circuit)
+    q_ma = _MeasurementQid('a', q0)
+    cirq.testing.assert_same_circuits(
+        deferred,
+        cirq.Circuit(
+            cirq.CX(q0, q_ma),
+            cirq.CX(q_ma, q1),
+            cirq.measure(q_ma, key='a'),
+            cirq.measure(q1, key='b'),
+        ),
+    )
+
+
+def test_sympy_control_complex():
+    q0, q1, q2 = cirq.LineQubit.range(3)
+    circuit = cirq.Circuit(
+        cirq.measure(q0, key='a'),
+        cirq.measure(q1, key='b'),
+        cirq.X(q2).with_classical_controls(sympy_parser.parse_expr('a >= b')),
+        cirq.measure(q2, key='c'),
+    )
+    assert_equivalent_to_deferred(circuit)
+    deferred = cirq.defer_measurements(circuit)
+    q_ma = _MeasurementQid('a', q0)
+    q_mb = _MeasurementQid('b', q1)
+    cirq.testing.assert_same_circuits(
+        deferred,
+        cirq.Circuit(
+            cirq.CX(q0, q_ma),
+            cirq.CX(q1, q_mb),
+            cirq.ControlledOperation(
+                [q_ma, q_mb], cirq.X(q2), cirq.SumOfProducts([[0, 0], [1, 0], [1, 1]])
+            ),
+            cirq.measure(q_ma, key='a'),
+            cirq.measure(q_mb, key='b'),
+            cirq.measure(q2, key='c'),
+        ),
+    )
+
+
+def test_sympy_control_multiqubit():
+    q0, q1, q2 = cirq.LineQubit.range(3)
+    circuit = cirq.Circuit(
+        cirq.measure(q0, q1, key='a'),
+        cirq.X(q2).with_classical_controls(sympy_parser.parse_expr('a >= 2')),
+        cirq.measure(q2, key='c'),
+    )
+    assert_equivalent_to_deferred(circuit)
+    deferred = cirq.defer_measurements(circuit)
+    q_ma0 = _MeasurementQid('a', q0)
+    q_ma1 = _MeasurementQid('a', q1)
+    cirq.testing.assert_same_circuits(
+        deferred,
+        cirq.Circuit(
+            cirq.CX(q0, q_ma0),
+            cirq.CX(q1, q_ma1),
+            cirq.ControlledOperation(
+                [q_ma0, q_ma1], cirq.X(q2), cirq.SumOfProducts([[1, 0], [1, 1]])
+            ),
+            cirq.measure(q_ma0, q_ma1, key='a'),
+            cirq.measure(q2, key='c'),
+        ),
+    )
+
+
 def test_nocompile_context():
     q0, q1 = cirq.LineQubit.range(2)
     circuit = cirq.Circuit(
@@ -295,15 +369,6 @@ def test_repr(qid: _MeasurementQid):
     test_repr(_MeasurementQid('0:1:a', cirq.LineQid(9, 4)))
 
 
-def test_sympy_control():
-    q0, q1 = cirq.LineQubit.range(2)
-    circuit = cirq.Circuit(
-        cirq.measure(q0, q1, key='a'), cirq.X(q1).with_classical_controls(sympy.Symbol('a'))
-    )
-    with pytest.raises(ValueError, match='Only KeyConditions are allowed'):
-        _ = cirq.defer_measurements(circuit)
-
-
 def test_confusion_map():
     q0, q1 = cirq.LineQubit.range(2)
     circuit = cirq.Circuit(