From e1aa6a4f8c3a5846b81913cd8aece6487d6bcb9d Mon Sep 17 00:00:00 2001 From: JerryChen97 Date: Tue, 20 Oct 2026 15:19:12 -0400 Subject: [PATCH 1/9] Support control-value operands in captured callables [sc-133997] Assisted-by: Codex --- doc/releases/changelog-dev.md | 4 +++ .../catalyst/from_plxpr/qfunc_interpreter.py | 11 ++++--- .../from_plxpr/qref_jax_primitives.py | 9 +++--- .../test/lit/test_qref/test_flat_circuits.py | 18 +++++++++++ .../from_plxpr/test_capture_integration.py | 30 +++++++++++++++++++ .../test/pytest/from_plxpr/test_from_plxpr.py | 19 +++++++----- 6 files changed, 76 insertions(+), 15 deletions(-) diff --git a/doc/releases/changelog-dev.md b/doc/releases/changelog-dev.md index 81bc4f6ff9..cdbc2165b4 100644 --- a/doc/releases/changelog-dev.md +++ b/doc/releases/changelog-dev.md @@ -808,6 +808,10 @@

Bug fixes ๐Ÿ›

+* Fixes mixed static and dynamic `control_values` for callable `qp.ctrl` under + `qjit(capture=True)`. + [(#XXXX)](https://github.com/PennyLaneAI/catalyst/pull/XXXX) + * MLIR passes that recreate `scf.for`, `scf.while`, `scf.if`, and `scf.index_switch` ops (`ctrl-lowering`, `adjoint-lowering`, `--convert-to-value-semantics`) now preserve discardable attributes such as `catalyst.estimated_iterations`, `catalyst.estimated_probability`, and diff --git a/frontend/catalyst/from_plxpr/qfunc_interpreter.py b/frontend/catalyst/from_plxpr/qfunc_interpreter.py index fb20641dcc..a45b0aaf7e 100644 --- a/frontend/catalyst/from_plxpr/qfunc_interpreter.py +++ b/frontend/catalyst/from_plxpr/qfunc_interpreter.py @@ -890,11 +890,14 @@ def handle_adjoint_transform(self, *invals, jaxpr, lazy, n_consts): # pylint: disable=unused-argument @PLxPRToQuantumJaxprInterpreter.register_primitive(plxpr_ctrl_transform_prim) -def handle_ctrl_transform(self, *invals, jaxpr, n_control, control_values, work_wires, n_consts): +def handle_ctrl_transform(self, *invals, jaxpr, n_control, work_wires, n_consts): """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 @@ -908,9 +911,9 @@ def handle_ctrl_transform(self, *invals, jaxpr, n_control, control_values, work_ *body.consts, *args, *control_qubits, + *control_values, n_control=n_control, jaxpr=body.jaxpr, - control_values=control_values, work_wires=work_wires, n_consts=len(body.consts), ) diff --git a/frontend/catalyst/from_plxpr/qref_jax_primitives.py b/frontend/catalyst/from_plxpr/qref_jax_primitives.py index a78ff8a595..4bb4043567 100644 --- a/frontend/catalyst/from_plxpr/qref_jax_primitives.py +++ b/frontend/catalyst/from_plxpr/qref_jax_primitives.py @@ -742,17 +742,18 @@ def _pl_ctrl_lowering( *plxpr_invals, jaxpr, n_control, - control_values, work_wires, n_consts, ): """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) diff --git a/frontend/test/lit/test_qref/test_flat_circuits.py b/frontend/test/lit/test_qref/test_flat_circuits.py index 773d7c8a0b..fe7dbc294b 100644 --- a/frontend/test/lit/test_qref/test_flat_circuits.py +++ b/frontend/test/lit/test_qref/test_flat_circuits.py @@ -75,6 +75,24 @@ def test_custom_op(i: int): print(test_custom_op.mlir) +# CHECK-LABEL: func.func public @test_dynamic_control_values( +# CHECK-SAME: %arg0: tensor, %arg1: tensor +@qp.qjit(capture=True, target="mlir", skip_preprocess=True, 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.""" + # CHECK: [[control_bool:%.+]] = stablehlo.compare NE, %arg1, {{%.+}}, UNSIGNED + # CHECK: [[control_tensor:%.+]] = stablehlo.convert [[control_bool]] : tensor + # CHECK: [[control_value:%.+]] = tensor.extract [[control_tensor]][] : tensor + # 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)) diff --git a/frontend/test/pytest/from_plxpr/test_capture_integration.py b/frontend/test/pytest/from_plxpr/test_capture_integration.py index 07461e291b..79896b968c 100644 --- a/frontend/test/pytest/from_plxpr/test_capture_integration.py +++ b/frontend/test/pytest/from_plxpr/test_capture_integration.py @@ -1503,6 +1503,36 @@ def c(x): assert qp.math.allclose(res, expected) +@pytest.mark.parametrize( + "operator_instance, container, static_first", + ((True, list, True), (True, list, False), (False, tuple, False), (False, list, False)), +) +def test_mixed_control_values(operator_instance, container, static_first): + """The four reported mixed-control cases follow the runtime Boolean value.""" + + def circuit(x, val): + qp.X(0 if static_first else 1) + values = container((True, val) if static_first else (val, True)) + if operator_instance: + qp.ctrl(qp.RX(x, 2), (0, 1), values) + else: + qp.ctrl(qp.RX, (0, 1), values)(x, 2) + return qp.expval(qp.Z(2)) + + compiled = qjit( + qp.QNode(circuit, qp.device("lightning.qubit", wires=4)), + capture=True, + skip_preprocess=True, + ) + reference = qp.QNode(circuit, qp.device("default.qubit", wires=4)) + x = jnp.array(0.5) + + for val in (False, True): + with qp.capture.pause(): + expected = reference(x, val) + assert qp.math.allclose(compiled(x, val), expected) + + @pytest.mark.parametrize("separate_funcs", (True, False)) def test_ctrl_transform_integration(separate_funcs): """Test that the ctrl transform can be applied.""" diff --git a/frontend/test/pytest/from_plxpr/test_from_plxpr.py b/frontend/test/pytest/from_plxpr/test_from_plxpr.py index 79eee21ed1..e322174154 100644 --- a/frontend/test/pytest/from_plxpr/test_from_plxpr.py +++ b/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 @@ -544,11 +545,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 @@ -636,8 +639,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 @@ -689,7 +693,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"] From 1217464afd1228503945c5cf6378f0881836a09c Mon Sep 17 00:00:00 2001 From: JerryChen97 Date: Tue, 20 Oct 2026 15:20:47 -0400 Subject: [PATCH 2/9] Link dynamic control changelog to #3373 [sc-133997] Assisted-by: Codex --- doc/releases/changelog-dev.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/doc/releases/changelog-dev.md b/doc/releases/changelog-dev.md index cdbc2165b4..b3fe05868e 100644 --- a/doc/releases/changelog-dev.md +++ b/doc/releases/changelog-dev.md @@ -810,7 +810,7 @@ * Fixes mixed static and dynamic `control_values` for callable `qp.ctrl` under `qjit(capture=True)`. - [(#XXXX)](https://github.com/PennyLaneAI/catalyst/pull/XXXX) + [(#3373)](https://github.com/PennyLaneAI/catalyst/pull/3373) * MLIR passes that recreate `scf.for`, `scf.while`, `scf.if`, and `scf.index_switch` ops (`ctrl-lowering`, `adjoint-lowering`, `--convert-to-value-semantics`) now preserve discardable From 893a7d84f0e20e98097eed56f7b286562a117a3b Mon Sep 17 00:00:00 2001 From: paul0403 Date: Wed, 21 Oct 2026 10:22:22 -0400 Subject: [PATCH 3/9] . --- .pre-commit-config.yaml | 30 +++++++++++++++--------------- 1 file changed, 15 insertions(+), 15 deletions(-) diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index 7befbda8e5..4e5fc6ecb9 100644 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -26,21 +26,21 @@ repos: "./pyproject.toml", ] -# - repo: local -# hooks: -# - id: pylint -# name: "๐Ÿ pylint ยท Lint frontend" -# entry: pylint -# language: system -# types: [python] -# args: -# [ -# "--errors-only", # match CI: pylint --errors-only frontend -# "-rn", # only display messages -# "-sn", # don't display the score -# "--rcfile=.pylintrc", # link to the repo config -# ] -# files: ^frontend/ +- repo: local + hooks: + - id: pylint + name: "๐Ÿ pylint ยท Lint frontend" + entry: pylint + language: system + types: [python] + args: + [ + "--errors-only", # match CI: pylint --errors-only frontend + "-rn", # only display messages + "-sn", # don't display the score + "--rcfile=.pylintrc", # link to the repo config + ] + files: ^frontend/ - repo: https://github.com/pre-commit/mirrors-clang-format rev: v20.1.8 From 4517b913acf002b9f2ccd8818cf705c649cfde3e Mon Sep 17 00:00:00 2001 From: paul0403 Date: Wed, 21 Oct 2026 10:27:49 -0400 Subject: [PATCH 4/9] .. --- frontend/test/lit/test_qref/test_flat_circuits.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/frontend/test/lit/test_qref/test_flat_circuits.py b/frontend/test/lit/test_qref/test_flat_circuits.py index fe7dbc294b..f312ae7342 100644 --- a/frontend/test/lit/test_qref/test_flat_circuits.py +++ b/frontend/test/lit/test_qref/test_flat_circuits.py @@ -77,7 +77,7 @@ def test_custom_op(i: int): # CHECK-LABEL: func.func public @test_dynamic_control_values( # CHECK-SAME: %arg0: tensor, %arg1: tensor -@qp.qjit(capture=True, target="mlir", skip_preprocess=True, collect_decomp_rules=False) +@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.""" From e251bbd7d7372090fc5ad10380c902e1d7a26e76 Mon Sep 17 00:00:00 2001 From: paul0403 Date: Wed, 21 Oct 2026 11:03:20 -0400 Subject: [PATCH 5/9] lit test check stablehlo better and add comment --- frontend/test/lit/test_qref/test_flat_circuits.py | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/frontend/test/lit/test_qref/test_flat_circuits.py b/frontend/test/lit/test_qref/test_flat_circuits.py index f312ae7342..c4b1976bc2 100644 --- a/frontend/test/lit/test_qref/test_flat_circuits.py +++ b/frontend/test/lit/test_qref/test_flat_circuits.py @@ -81,7 +81,11 @@ def test_custom_op(i: int): @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.""" - # CHECK: [[control_bool:%.+]] = stablehlo.compare NE, %arg1, {{%.+}}, UNSIGNED + # COM: stablehlo compares bools against "not false" due to overly conservative lowering of + # COM: the convert_element_type primitive + # CHECK: [[f:%.+]] = stablehlo.constant dense : tensor + # CHECK: [[broadcast:%.+]] = stablehlo.broadcast_in_dim [[f]], dims = [] : (tensor) -> tensor + # CHECK: [[control_bool:%.+]] = stablehlo.compare NE, %arg1, [[broadcast]], UNSIGNED # CHECK: [[control_tensor:%.+]] = stablehlo.convert [[control_bool]] : tensor # CHECK: [[control_value:%.+]] = tensor.extract [[control_tensor]][] : tensor # CHECK: qref.ctrl({{[^)]+}}) ctrlvals({{[^,]+}}, [[control_value]]) { From 16a0fdc1d5f1d7ac3ccf2813608df768be46291d Mon Sep 17 00:00:00 2001 From: JerryChen97 Date: Wed, 21 Oct 2026 11:35:07 -0400 Subject: [PATCH 6/9] Clarify static control wire in mixed-control test Assisted-by: Codex --- .../pytest/from_plxpr/test_capture_integration.py | 12 +++++++----- 1 file changed, 7 insertions(+), 5 deletions(-) diff --git a/frontend/test/pytest/from_plxpr/test_capture_integration.py b/frontend/test/pytest/from_plxpr/test_capture_integration.py index 79896b968c..3659a38c72 100644 --- a/frontend/test/pytest/from_plxpr/test_capture_integration.py +++ b/frontend/test/pytest/from_plxpr/test_capture_integration.py @@ -1504,15 +1504,17 @@ def c(x): @pytest.mark.parametrize( - "operator_instance, container, static_first", - ((True, list, True), (True, list, False), (False, tuple, False), (False, list, False)), + "operator_instance, container, static_control_wire", + ((True, list, 0), (True, list, 1), (False, tuple, 1), (False, list, 1)), ) -def test_mixed_control_values(operator_instance, container, static_first): +def test_mixed_control_values(operator_instance, container, static_control_wire): """The four reported mixed-control cases follow the runtime Boolean value.""" def circuit(x, val): - qp.X(0 if static_first else 1) - values = container((True, val) if static_first else (val, True)) + qp.X(static_control_wire) + values = [val, val] + values[static_control_wire] = True + values = container(values) if operator_instance: qp.ctrl(qp.RX(x, 2), (0, 1), values) else: From 70920e1fe237942625f80fda9d4029ff909b6db6 Mon Sep 17 00:00:00 2001 From: JerryChen97 Date: Wed, 21 Oct 2026 11:40:30 -0400 Subject: [PATCH 7/9] Bump PennyLane to 0.46.0-dev132 Assisted-by: Codex --- .dep-versions | 2 +- doc/requirements.txt | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/.dep-versions b/.dep-versions index dc5d78feeb..b7a2aea19c 100644 --- a/.dep-versions +++ b/.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-dev132 # For a custom LQ/LK version, update the package version here and at # 'doc/requirements.txt' diff --git a/doc/requirements.txt b/doc/requirements.txt index 38fa1ce790..c30421175c 100644 --- a/doc/requirements.txt +++ b/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-dev132 From c4994ba156a1ed702b0bdf4d6205aa65a90b28e8 Mon Sep 17 00:00:00 2001 From: "Yushao Chen (Jerry)" Date: Wed, 21 Oct 2026 13:21:56 -0400 Subject: [PATCH 8/9] Bump to 133 Co-authored-by: Yushao Chen (Jerry) --- .dep-versions | 2 +- doc/requirements.txt | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/.dep-versions b/.dep-versions index b7a2aea19c..80e65ff138 100644 --- a/.dep-versions +++ b/.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-dev132 +pennylane=0.46.0-dev133 # For a custom LQ/LK version, update the package version here and at # 'doc/requirements.txt' diff --git a/doc/requirements.txt b/doc/requirements.txt index c30421175c..3d70a87b31 100644 --- a/doc/requirements.txt +++ b/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-dev132 +pennylane==0.46.0-dev133 From 641c2c08a7f2b8c0cb3d1c25e42a74d19d0c6b1b Mon Sep 17 00:00:00 2001 From: JerryChen97 Date: Wed, 21 Oct 2026 14:53:19 -0400 Subject: [PATCH 9/9] human readible tests with reasonable logic and structure, aligned with how we test usually --- .../from_plxpr/test_capture_integration.py | 52 ++++++++++--------- 1 file changed, 27 insertions(+), 25 deletions(-) diff --git a/frontend/test/pytest/from_plxpr/test_capture_integration.py b/frontend/test/pytest/from_plxpr/test_capture_integration.py index 3659a38c72..1b3a0ed14e 100644 --- a/frontend/test/pytest/from_plxpr/test_capture_integration.py +++ b/frontend/test/pytest/from_plxpr/test_capture_integration.py @@ -1503,36 +1503,38 @@ def c(x): assert qp.math.allclose(res, expected) -@pytest.mark.parametrize( - "operator_instance, container, static_control_wire", - ((True, list, 0), (True, list, 1), (False, tuple, 1), (False, list, 1)), -) -def test_mixed_control_values(operator_instance, container, static_control_wire): - """The four reported mixed-control cases follow the runtime Boolean value.""" +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(static_control_wire) - values = [val, val] - values[static_control_wire] = True - values = container(values) - if operator_instance: - qp.ctrl(qp.RX(x, 2), (0, 1), values) - else: - qp.ctrl(qp.RX, (0, 1), values)(x, 2) - return qp.expval(qp.Z(2)) + 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)) - compiled = qjit( - qp.QNode(circuit, qp.device("lightning.qubit", wires=4)), - capture=True, - skip_preprocess=True, - ) - reference = qp.QNode(circuit, qp.device("default.qubit", wires=4)) 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.""" - for val in (False, True): - with qp.capture.pause(): - expected = reference(x, val) - assert qp.math.allclose(compiled(x, val), expected) + @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))