Support dynamic control values in captured qp.ctrl

.dep-versions

@@ -8,7 +8,7 @@ enzyme=v0.0.238
 
 # For a custom PL version, update the package version here and at
 # 'doc/requirements.txt'
-pennylane=0.46.0-dev131
+pennylane=0.46.0-dev133
 
 # For a custom LQ/LK version, update the package version here and at
 # 'doc/requirements.txt'

doc/releases/changelog-dev.md

@@ -768,7 +768,7 @@
 * The `--convert-to-value-semantics` and `--convert-to-reference-semantics` passes now collect
   subroutine symbol uses only once instead of once per subroutine, improving their performance.
   [(#3341)](https://github.com/PennyLaneAI/catalyst/pull/3341)
-  
+
 * Add `post-device-pipeline` to the `transform-sequence` when `qp.qjit(capture=True)`
   [(#3312)](https://github.com/PennyLaneAI/catalyst/pull/3312)
 
@@ -811,6 +811,10 @@
 
 <h3>Bug fixes 🐛</h3>
 
+* Fixes mixed static and dynamic `control_values` for callable `qp.ctrl` under
+  `qjit(capture=True)`.
+  [(#3373)](https://github.com/PennyLaneAI/catalyst/pull/3373)
+
 * Controlled quantum regions now keeps track of the `work_wire_type` passed from `qp.ctrl`.
   [(#3366)](https://github.com/PennyLaneAI/catalyst/pull/3366)
 

doc/requirements.txt

@@ -34,4 +34,4 @@ lxml_html_clean
 --extra-index-url https://test.pypi.org/simple/
 pennylane-lightning-kokkos==0.46.0-dev27
 pennylane-lightning==0.46.0-dev27
-pennylane==0.46.0-dev131
+pennylane==0.46.0-dev133

frontend/catalyst/from_plxpr/qfunc_interpreter.py

@@ -891,19 +891,15 @@ def handle_adjoint_transform(self, *invals, jaxpr, lazy, n_consts):
 # pylint: disable=unused-argument,too-many-arguments
 @PLxPRToQuantumJaxprInterpreter.register_primitive(plxpr_ctrl_transform_prim)
 def handle_ctrl_transform(
-    self,
-    *invals,
-    jaxpr,
-    n_control,
-    control_values,
-    work_wires,
-    n_consts,
-    work_wire_type="borrowed",
+    self, *invals, jaxpr, n_control, work_wires, n_consts, work_wire_type="borrowed"
 ):
     """Interpret a control transform, then re-bind it for lowering to a `qref.ctrl` region op."""
+    values_start = len(invals) - n_control
+    controls_start = values_start - n_control
     consts = invals[:n_consts]
-    args = invals[n_consts:-n_control]
-    control_wires = invals[-n_control:]
+    args = invals[n_consts:controls_start]
+    control_wires = invals[controls_start:values_start]
+    control_values = invals[values_start:]
 
     control_qubits = [
         w if is_abstract_qubit(w) else qref_get_p.bind(self.init_qreg, w) for w in control_wires
@@ -917,9 +913,9 @@ def handle_ctrl_transform(
         *body.consts,
         *args,
         *control_qubits,
+        *control_values,
         n_control=n_control,
         jaxpr=body.jaxpr,
-        control_values=control_values,
         work_wires=work_wires,
         work_wire_type=work_wire_type,
         n_consts=len(body.consts),

frontend/catalyst/from_plxpr/qref_jax_primitives.py

@@ -742,18 +742,19 @@ def _pl_ctrl_lowering(
     *plxpr_invals,
     jaxpr,
     n_control,
-    control_values,
     work_wires,
     n_consts,
     work_wire_type="borrowed",
 ):
     """Lower a `qp.ctrl` transform to a `qref.ctrl` region op."""
-    body_invals = plxpr_invals[: len(plxpr_invals) - n_control]
-    control_qubits = list(plxpr_invals[len(plxpr_invals) - n_control :])
+    values_start = len(plxpr_invals) - n_control
+    controls_start = values_start - n_control
+    body_invals = plxpr_invals[:controls_start]
+    control_qubits = list(plxpr_invals[controls_start:values_start])
 
     i1 = ir.IntegerType.get_signless(1)
     control_values_i1 = [
-        ConstantOp(i1, ir.IntegerAttr.get(i1, 1 if v else 0)).result for v in control_values
+        TensorExtractOp(i1, value, []).result for value in plxpr_invals[values_start:]
     ]
 
     op = CtrlOp(ctrl_qubits=control_qubits, ctrl_values=control_values_i1)

frontend/test/lit/test_qref/test_flat_circuits.py

@@ -75,6 +75,28 @@ def test_custom_op(i: int):
 print(test_custom_op.mlir)
 
 
+# CHECK-LABEL: func.func public @test_dynamic_control_values(
+# CHECK-SAME: %arg0: tensor<f64>, %arg1: tensor<i1>
+@qp.qjit(capture=True, target="mlir", collect_decomp_rules=False)
+@qp.qnode(qp.device("null.qubit", wires=3))
+def test_dynamic_control_values(x: float, val: bool):
+    """Dynamic control operands lower to SSA values in the controlled region."""
+    # COM: stablehlo compares bools against "not false" due to overly conservative lowering of
+    # COM: the convert_element_type primitive
+    # CHECK: [[f:%.+]] = stablehlo.constant dense<false> : tensor<i1>
+    # CHECK: [[broadcast:%.+]] = stablehlo.broadcast_in_dim [[f]], dims = [] : (tensor<i1>) -> tensor<i1>
+    # CHECK: [[control_bool:%.+]] = stablehlo.compare NE, %arg1, [[broadcast]], UNSIGNED
+    # CHECK: [[control_tensor:%.+]] = stablehlo.convert [[control_bool]] : tensor<i1>
+    # CHECK: [[control_value:%.+]] = tensor.extract [[control_tensor]][] : tensor<i1>
+    # CHECK: qref.ctrl({{[^)]+}}) ctrlvals({{[^,]+}}, [[control_value]]) {
+    # CHECK: qref.custom "RX"
+    qp.ctrl(qp.RX, (0, 1), [True, val])(x, 2)
+    return qp.expval(qp.Z(2))
+
+
+print(test_dynamic_control_values.mlir)
+
+
 # CHECK: func.func public @test_measure() -> tensor<2xf64>
 @qp.qjit(capture=True, target="mlir", collect_decomp_rules=False)
 @qp.qnode(qp.device("null.qubit", wires=3))

frontend/test/pytest/from_plxpr/test_capture_integration.py

@@ -1503,6 +1503,40 @@ def c(x):
     assert qp.math.allclose(res, expected)
 
 
+def test_mixed_control_values_operator(backend):
+    """Test both orders of mixed static and dynamic controls on RX instances."""
+
+    @qjit(capture=True, skip_preprocess=True)
+    @qp.qnode(qp.device(backend, wires=4))
+    def circuit(x, val):
+        qp.X(0)
+        qp.X(1)
+        qp.ctrl(qp.RX(x, 2), (0, 1), [True, val])
+        qp.ctrl(qp.RX(x, 3), (0, 1), [val, True])
+        return qp.expval(qp.Z(2)), qp.expval(qp.Z(3))
+
+    x = jnp.array(0.5)
+    assert qp.math.allclose(circuit(x, False), (1.0, 1.0))
+    assert qp.math.allclose(circuit(x, True), (jnp.cos(x), jnp.cos(x)))
+
+
+def test_mixed_control_values_callable(backend):
+    """Test tuple and list controls mixing static and dynamic values on the RX callable."""
+
+    @qjit(capture=True, skip_preprocess=True)
+    @qp.qnode(qp.device(backend, wires=4))
+    def circuit(x, val):
+        qp.X(0)
+        qp.X(1)
+        qp.ctrl(qp.RX, (0, 1), (val, True))(x, 2)
+        qp.ctrl(qp.RX, (0, 1), [val, True])(x, 3)
+        return qp.expval(qp.Z(2)), qp.expval(qp.Z(3))
+
+    x = jnp.array(0.5)
+    assert qp.math.allclose(circuit(x, False), (1.0, 1.0))
+    assert qp.math.allclose(circuit(x, True), (jnp.cos(x), jnp.cos(x)))
+
+
 @pytest.mark.parametrize("separate_funcs", (True, False))
 def test_ctrl_transform_integration(separate_funcs):
     """Test that the ctrl transform can be applied."""

frontend/test/pytest/from_plxpr/test_from_plxpr.py

@@ -513,8 +513,9 @@ def circuit():
         eqn = qfunc_xpr.eqns[3]
         assert eqn.primitive == ctrl_transform_prim
         assert eqn.params["n_control"] == 1
-        assert eqn.params["control_values"] == (False,)
-        assert eqn.invars[-1] is qfunc_xpr.eqns[2].outvars[0]
+        assert "control_values" not in eqn.params
+        assert eqn.invars[-2] is qfunc_xpr.eqns[2].outvars[0]
+        assert not eqn.invars[-1].val
         body_xpr = eqn.params["jaxpr"]
         s_eqn = body_xpr.eqns[1]
         assert s_eqn.primitive == qref_operator_p
@@ -562,11 +563,13 @@ def c():
             outer = qfunc_xpr.eqns[3]
             assert outer.primitive == ctrl_transform_prim
             assert outer.params["n_control"] == 1
-            assert outer.params["control_values"] == (False,)
+            assert "control_values" not in outer.params
+            assert not outer.invars[-1].val
             inner = outer.params["jaxpr"].eqns[1]
             assert inner.primitive == ctrl_transform_prim
             assert inner.params["n_control"] == 1
-            assert inner.params["control_values"] == (True,)
+            assert "control_values" not in inner.params
+            assert inner.invars[-1].val
             s_eqn = inner.params["jaxpr"].eqns[1]
             assert s_eqn.primitive == qref_operator_p
             assert s_eqn.params["n_ctrls"] == 0
@@ -654,8 +657,9 @@ def c(wire):
             eqn = qfunc_xpr.eqns[6]
             assert eqn.primitive == ctrl_transform_prim
             assert eqn.params["n_control"] == 1
-            assert eqn.params["control_values"] == (True,)
-            assert eqn.invars[-1] is qfunc_xpr.eqns[5].outvars[0]
+            assert "control_values" not in eqn.params
+            assert eqn.invars[-2] is qfunc_xpr.eqns[5].outvars[0]
+            assert eqn.invars[-1].val
             t_eqn = eqn.params["jaxpr"].eqns[1]
             assert t_eqn.primitive == qref_operator_p
             assert t_eqn.params["n_ctrls"] == 0
@@ -707,7 +711,8 @@ def c():
         ctrl_eqn = qfunc_xpr.eqns[4]
         assert ctrl_eqn.primitive == ctrl_transform_prim
         assert ctrl_eqn.params["n_control"] == 2
-        assert ctrl_eqn.params["control_values"] == (True, True)
+        assert "control_values" not in ctrl_eqn.params
+        assert [var.val for var in ctrl_eqn.invars[-2:]] == [True, True]
 
         for_loop_eqn = ctrl_eqn.params["jaxpr"].eqns[0]
         for_loop_xpr = for_loop_eqn.params["jaxpr_body_fn"]