Simplify controlled gate for SumOfProducts (quantumlib#5873)

Co-authored-by: Cirq Bot <[email protected]>

Author: daxfohl

cirq/ops/common_gates_test.py

@@ -125,6 +125,7 @@ def test_specialized_control(input_gate, specialized_output):
     assert input_gate.controlled() == specialized_output
     assert input_gate.controlled(num_controls=1) == specialized_output
     assert input_gate.controlled(control_values=((1,),)) == specialized_output
+    assert input_gate.controlled(control_values=cirq.SumOfProducts([[1]])) == specialized_output
     assert input_gate.controlled(control_qid_shape=(2,)) == specialized_output
     assert np.allclose(
         cirq.unitary(specialized_output),
@@ -166,6 +167,28 @@ def test_specialized_control(input_gate, specialized_output):
     )
 
 
+@pytest.mark.parametrize(
+    'input_gate, specialized_output',
+    [
+        (cirq.Z, cirq.CCZ),
+        (cirq.X, cirq.CCX),
+        (cirq.ZPowGate(exponent=0.5), cirq.CCZPowGate(exponent=0.5)),
+        (cirq.XPowGate(exponent=0.5), cirq.CCXPowGate(exponent=0.5)),
+    ],
+)
+def test_specialized_control_two_step(input_gate, specialized_output):
+    # Two-qubit control on the input gate gives the specialized output
+    assert input_gate.controlled().controlled() == specialized_output
+    assert input_gate.controlled(num_controls=2) == specialized_output
+    assert input_gate.controlled(control_values=[1, 1]) == specialized_output
+    assert input_gate.controlled(control_values=cirq.SumOfProducts([[1, 1]])) == specialized_output
+    assert input_gate.controlled(control_qid_shape=(2, 2)) == specialized_output
+    assert np.allclose(
+        cirq.unitary(specialized_output),
+        cirq.unitary(cirq.ControlledGate(input_gate, num_controls=2)),
+    )
+
+
 @pytest.mark.parametrize(
     'gate, specialized_type',
     [

cirq/ops/controlled_gate.py

@@ -80,6 +80,11 @@ def __init__(
                 bounds, or if the sub_gate is not a unitary or mixture.
         """
         _validate_sub_object(sub_gate)
+
+        # Simplify a single SumOfProducts
+        if isinstance(control_values, cv.SumOfProducts) and len(control_values._conjunctions) == 1:
+            control_values = control_values._conjunctions[0]
+
         if num_controls is None:
             if control_values is not None:
                 num_controls = (

cirq/transformers/measurement_transformers.py

@@ -121,12 +121,9 @@ def defer(op: 'cirq.Operation', _) -> 'cirq.OP_TREE':
                     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)
+                    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)
                 else:
                     raise ValueError('Only KeyConditions are allowed.')