Fix mixed control values in qjit capture

doc/releases/changelog-dev.md

@@ -1336,6 +1336,11 @@
 
 <h3>Breaking changes 💔</h3>
 
+* The `ctrl_transform` capture primitive now receives Boolean control values as operands
+  after the control wires. Its `control_values` parameter has been removed, so consuming
+  interpreters, including Catalyst, must support the new operand layout.
+  [(#10356)](https://github.com/PennyLaneAI/pennylane/pull/10356)
+
 * :func:`~.transforms.decompose` is now a pass-only transform that always dispatches to Catalyst's
   ``graph_decomposition`` under :func:`~.qjit`. The previous tape-based implementation is
   available as :func:`~.transforms._tape_decompose`.
@@ -2056,6 +2061,10 @@
 
 <h3>Bug fixes 🐛</h3>
 
+* :func:`~.ctrl` now supports mixed static and dynamic control values under
+  `qp.qjit(capture=True)`, for both operator instances and callables.
+  [(#10356)](https://github.com/PennyLaneAI/pennylane/pull/10356)
+
 * `qp.cond` can now handle `Operator2` instances being provided instead of branch functions.
   While this syntax was supported with tape-based PennyLane, it had never been
   properly supported with program capture and Catalyst.

pennylane/capture/base_interpreter.py

@@ -487,19 +487,19 @@ def handle_adjoint_transform(self, *invals, jaxpr, lazy, n_consts):
 
 # pylint: disable=too-many-arguments
 @PlxprInterpreter.register_primitive(ctrl_transform_prim)
-def handle_ctrl_transform(self, *invals, n_control, jaxpr, control_values, work_wires, n_consts):
+def handle_ctrl_transform(self, *invals, n_control, jaxpr, work_wires, n_consts):
     """Interpret a ctrl transform primitive."""
+    controls_start = len(invals) - 2 * n_control
     consts = invals[:n_consts]
-    args = invals[n_consts:-n_control]
+    args = invals[n_consts:controls_start]
     jaxpr = jaxpr_to_jaxpr(copy(self), jaxpr, consts, *args)
 
     return ctrl_transform_prim.bind(
         *jaxpr.consts,
         *args,
-        *invals[-n_control:],
+        *invals[controls_start:],
         n_control=n_control,
         jaxpr=jaxpr.jaxpr,
-        control_values=control_values,
         work_wires=work_wires,
         n_consts=len(jaxpr.consts),
     )

pennylane/ops/op_math/controlled.py

@@ -300,7 +300,7 @@ def create_controlled_op(op, control, control_values, work_wires, work_wire_type
     if isinstance(control_values, (int, bool)):
         control_values = [bool(control_values)]
     elif isinstance(control_values, tuple):
-        control_values = [bool(v) for v in control_values]
+        control_values = [v if math.is_abstract(v) else bool(v) for v in control_values]
 
     qp.QueuingManager.remove(op)
 
@@ -418,14 +418,17 @@ def wrapper(*args, **kwargs):
 
 
 @ctrl_transform_prim.def_impl
-def _ctrl_transform_impl(*args, n_control, jaxpr, control_values, work_wires, n_consts):
+def _ctrl_transform_impl(*args, n_control, jaxpr, work_wires, n_consts):
     from pennylane.tape.plxpr_conversion import (  # pylint: disable=import-outside-toplevel
         CollectOpsandMeas,
     )
 
+    values_start = len(args) - n_control
+    controls_start = values_start - n_control
     consts = args[:n_consts]
-    control_wires = args[-n_control:]
-    args = args[n_consts:-n_control]
+    control_wires = args[controls_start:values_start]
+    control_values = args[values_start:]
+    args = args[n_consts:controls_start]
 
     collector = CollectOpsandMeas()
     with qp.QueuingManager.stop_recording():
@@ -458,9 +461,9 @@ def new_qfunc(*args, **kwargs):
             *abstract_shapes,
             *flat_args,
             *control_wires,
+            *(math.cast(value, dtype=bool) for value in control_values),
             jaxpr=jaxpr.jaxpr,
             n_control=len(control_wires),
-            control_values=control_values,
             work_wires=qp.wires.Wires(work_wires) if work_wires is not None else work_wires,
             n_consts=len(jaxpr.consts),
         )

pennylane/ops/op_math/controlled2.py

@@ -67,7 +67,9 @@ def _setup_control_values(control_values, num_control_wires):
         raise ValueError("control_values should be the same length as control_wires")
 
     if isinstance(control_values, (list, tuple)):
-        control_values = qp.math.asarray(control_values, like=control_values[0])
+        control_values = qp.math.asarray(
+            control_values, like=qp.math.get_interface(*control_values)
+        )
 
     if not isinstance(control_values, AbstractArray):
         control_values = qp.math.cast(control_values, dtype=bool)

pennylane/tape/plxpr_conversion.py

@@ -131,17 +131,22 @@ def _ctrl_transform_prim(self, *invals, n_control, jaxpr, n_consts, **params):
     """Handle a control transform primitive by collecting the operations in the jaxpr,
     and then applying their controlled versions.
     """
+    values_start = len(invals) - n_control
+    controls_start = values_start - n_control
     consts = invals[:n_consts]
-    args = invals[n_consts:-n_control]
-    control = math.array(invals[-n_control:])
+    args = invals[n_consts:controls_start]
+    control = math.array(invals[controls_start:values_start])
+    control_values = invals[values_start:]
 
     child = CollectOpsandMeas()
     child.eval(jaxpr, consts, *args)
     assert child.state
 
     for op in child.state["ops"]:
         with pause():
-            self.state["ops"].append(ops.ctrl(op, control=control, **params))
+            self.state["ops"].append(
+                ops.ctrl(op, control=control, control_values=control_values, **params)
+            )
 
     return []
 

tests/capture/test_base_interpreter.py

@@ -391,15 +391,17 @@ def test_ctrl_transform(self):
         """Test the higher order ctrl transform."""
 
         @SimplifyInterpreter()
-        def f(x, control):
+        def f(x, control, val):
             def g(y):
                 _ = qp.RY(y, 0) ** 3
 
-            qp.ctrl(g, control)(x)
+            qp.ctrl(g, control, control_values=[val])(x)
 
-        jaxpr = jax.make_jaxpr(f)(0.5, 1)
+        jaxpr = jax.make_jaxpr(f)(0.5, 1, False)
 
         assert jaxpr.eqns[0].primitive == ctrl_transform_prim
+        assert jaxpr.eqns[0].invars[-2:] == jaxpr.jaxpr.invars[1:]
+        assert "control_values" not in jaxpr.eqns[0].params
         inner_jaxpr = jaxpr.eqns[0].params["jaxpr"]
         # first eqn mul, second RY
         assert_eqn_matches_op(inner_jaxpr.eqns[1], qp.RY)

tests/capture/test_nested_plxpr.py

@@ -324,7 +324,7 @@ def f(x, w):
         qp.assert_equal(q.queue[0], expected)
 
         assert plxpr.eqns[0].primitive == ctrl_transform_prim
-        assert plxpr.eqns[0].params["control_values"] == (True,)
+        assert plxpr.eqns[0].invars[-1].val
         assert plxpr.eqns[0].params["n_control"] == 1
         assert plxpr.eqns[0].params["work_wires"] is None
         assert plxpr.eqns[0].params["n_consts"] == 0
@@ -346,7 +346,7 @@ def f(w1, w2, w3):
         assert len(q) == 1
 
         assert plxpr.eqns[0].primitive == ctrl_transform_prim
-        assert plxpr.eqns[0].params["control_values"] == (True, True)
+        assert [value.val for value in plxpr.eqns[0].invars[-2:]] == [True, True]
         assert plxpr.eqns[0].params["n_control"] == 2
         assert plxpr.eqns[0].params["work_wires"] is None
 
@@ -384,9 +384,35 @@ def f(z):
         qp.assert_equal(q.queue[0], expected)
         assert len(q) == 1
 
-        assert plxpr.eqns[0].params["control_values"] == (False, True)
+        assert [value.val for value in plxpr.eqns[0].invars[-2:]] == [False, True]
         assert plxpr.eqns[0].params["n_control"] == 2
 
+    @pytest.mark.parametrize("container, static_first", ((tuple, False), (list, True)))
+    def test_mixed_control_values(self, container, static_first):
+        """Capture supplies runtime control-value operands for qjit lowering."""
+
+        def circuit(x, val):
+            values = (True, val) if static_first else (val, True)
+            qp.ctrl(qp.RX, (0, 1), container(values))(x, 2)
+
+        with jax.checking_leaks():
+            plxpr = jax.make_jaxpr(circuit)(0.5, False)
+
+        assert len(plxpr.eqns) == 1
+        ctrl_eqn = plxpr.eqns[0]
+        assert ctrl_eqn.primitive == ctrl_transform_prim
+        assert "control_values" not in ctrl_eqn.params
+        assert ctrl_eqn.params["n_control"] == 2
+        value_args = ctrl_eqn.invars[-2:]
+        assert value_args[int(static_first)] == plxpr.jaxpr.invars[1]
+        assert value_args[1 - int(static_first)].val
+        assert [wire.val for wire in ctrl_eqn.invars[-4:-2]] == [0, 1]
+
+        with qp.queuing.AnnotatedQueue() as queue:
+            jax.core.eval_jaxpr(plxpr.jaxpr, plxpr.consts, 1.2, False)
+        expected_values = [True, False] if static_first else [False, True]
+        qp.assert_equal(queue.queue[0], qp.ctrl(qp.RX(1.2, 2), (0, 1), expected_values))
+
     def test_nested_control(self):
         """Test that control can be nested."""
 
@@ -402,8 +428,8 @@ def f(x, w1, w2):
             plxpr.eqns[1].invars[0] is plxpr.jaxpr.invars[1]
         )  # first input is first control wire, const
         assert plxpr.eqns[1].invars[1] is plxpr.jaxpr.invars[0]  # second input is x, first arg
-        assert plxpr.eqns[1].invars[-1] is plxpr.jaxpr.invars[2]  # second control wire
-        assert len(plxpr.eqns[1].invars) == 6  # one const, 4 args, one control wire
+        assert plxpr.eqns[1].invars[-2] is plxpr.jaxpr.invars[2]  # second control wire
+        assert len(plxpr.eqns[1].invars) == 7  # one const, 4 args, one control wire and value
 
         with qp.queuing.AnnotatedQueue() as q:
             jax.core.eval_jaxpr(plxpr.jaxpr, plxpr.consts, 1.2, 3, 4)
@@ -439,7 +465,7 @@ def workflow(wire):
             qp.assert_equal(q.queue[2], qp.ctrl(qp.S(2), 0))
 
         eqn = jaxpr.eqns[0]
-        assert eqn.params["control_values"] == (True,)
+        assert eqn.invars[-1].val
         assert eqn.params["n_consts"] == 0
         assert eqn.params["n_control"] == 1
         assert eqn.params["work_wires"] is None

tests/ops/op_math/test_controlled2.py

@@ -636,6 +636,8 @@ def circ(cvals):
             qp.ctrl(op, control=[3, 4], control_values=[0, 1])
             op2 = qp.ctrl(qp.H(0), control=[1, 2], control_values=[1, 0])
             qp.ctrl(op2, control=[3, 4], control_values=cvals)
+            qp.ctrl(qp.H(0), control=[1, 2], control_values=[True, cvals[0]])
+            qp.ctrl(qp.H(0), control=[1, 2], control_values=[cvals[0], True])
 
         jaxpr = jax.make_jaxpr(circ)(jax.numpy.array([True, False]))
         tape = qp.testing.plxpr_to_tape(jaxpr.jaxpr, jaxpr.consts, jax.numpy.array([False, True]))
@@ -644,6 +646,8 @@ def circ(cvals):
             qp.ctrl(qp.H(1), control=[2, 3, 0], control_values=[0, 1, 1]),
             qp.ctrl(qp.H(0), control=[3, 4, 1, 2], control_values=[0, 1, 0, 1]),
             qp.ctrl(qp.H(0), control=[3, 4, 1, 2], control_values=[0, 1, 1, 0]),
+            qp.ctrl(qp.H(0), control=[1, 2], control_values=[1, 0]),
+            qp.ctrl(qp.H(0), control=[1, 2], control_values=[0, 1]),
         ]
 
     def test_old_decomp_integration(self):

tests/resource/test_specs.py

@@ -163,6 +163,17 @@ def circuit():
             ),
         )
 
+    @pytest.mark.xfail(
+        strict=True,
+        raises=pytest.RaisesExc(
+            TypeError,
+            match=(
+                r"^handle_ctrl_transform\(\) missing 1 required keyword-only "
+                r"argument: 'control_values'"
+            ),
+        ),
+        reason="Requires Catalyst control-value operand support (sc-133997).",
+    )
     def test_complex(self):
         """Test a complex case of qp.specs() against PennyLane"""
 

tests/tape/test_plxpr_conversion.py

@@ -227,12 +227,12 @@ def qfunc(x, wire):
             qp.RX(x, wire)
             qp.X(wire)
 
-        def f(x):
-            qp.ctrl(qfunc, control=[1, 2], control_values=[False, False])(x, 0)
+        def f(x, val):
+            qp.ctrl(qfunc, control=[1, 2], control_values=[False, val])(x, 0)
 
         obj = CollectOpsandMeas()
         x = jax.numpy.array(-0.98)
-        obj(f)(x)
+        obj(f)(x, False)
 
         assert len(obj.state["ops"]) == 2
         expected0 = qp.ctrl(qp.RX(x, 0), [1, 2], control_values=[False, False])