Repository: PennyLaneAI/pennylane
Pull request #7199: [Capture] User transforms are expanded before execution
Author: login: mudit2812
Created: 2025-04-03T20:02:27Z | Merged: 2025-04-04T21:02:45Z

[sc-83799]

**Context:**
This story adds support for expanding user transforms before sending jaxpr to devices when executing a QNode.

**Description of the Change:**
* Update `qnode_prim.impl` to expand any present user transforms in the `qfunc_jaxpr` before applying device preprocessing transforms, using `qml.capture.expand_plxpr_transforms`.
* Update `TransformDispatcher` to treat QNodes differently from other callables when program capture is enabled. When transforming a `QNode`, the transform will also be added to the `QNode`'s tranform program. This does not change the captured program, but it gives us the ability to use the transform program to apply qnode transforms to the `qfunc_jaxpr`.
* Update `transform_prim.impl` to just return the evaluation of the inner jaxpr. This was done so that calling a transformed qnode (which is actually a wrapper function around the qnode) will not raise a `NotImplementedError`.

**Benefits:**
Users no longer need to manually decorate their qnode with `@qml.capture.expand_plxpr_transforms`.

**Possible Drawbacks:**


**Related GitHub Issues:**



Diff:
diff --git a/doc/releases/changelog-dev.md b/doc/releases/changelog-dev.md
index ff64b9c3413..2edb4c27160 100644
--- a/doc/releases/changelog-dev.md
+++ b/doc/releases/changelog-dev.md
@@ -164,6 +164,10 @@ With `qml.decompositions.enable_graph()`, the following new features are availab
 
 <h4>Capturing and Representing Hybrid Programs 📥</h4>
 
+* Transformations can now be directly applied to a `QNode` with program capture enabled without having
+  to use the `@qml.capture.expand_plxpr_transforms` decorator.
+  [(#7199)](https://github.com/PennyLaneAI/pennylane/pull/7199)
+
 * Python control flow (`if/else`, `for`, `while`) is now supported when program capture is enabled by setting
   `autograph=True` at the QNode level.
   [(#6837)](https://github.com/PennyLaneAI/pennylane/pull/6837)
diff --git a/pennylane/transforms/core/transform_dispatcher.py b/pennylane/transforms/core/transform_dispatcher.py
index d71bec29fbd..79e35bd98d9 100644
--- a/pennylane/transforms/core/transform_dispatcher.py
+++ b/pennylane/transforms/core/transform_dispatcher.py
@@ -191,7 +191,8 @@ def processing_fn(results):
 
         if isinstance(obj, qml.QNode):
             if qml.capture.enabled():
-                return self._capture_callable_transform(obj, targs, tkwargs)
+                new_qnode = self.default_qnode_transform(obj, targs, tkwargs)
+                return self._capture_callable_transform(new_qnode, targs, tkwargs)
             return self._qnode_transform(obj, targs, tkwargs)
 
         if isinstance(obj, qml.devices.Device):
@@ -585,7 +586,9 @@ def _create_transform_primitive(name):
     def _(
         *all_args, inner_jaxpr, args_slice, consts_slice, targs_slice, tkwargs
     ):  # pylint: disable=unused-argument
-        raise NotImplementedError
+        args = all_args[args_slice]
+        consts = all_args[consts_slice]
+        return qml.capture.eval_jaxpr(inner_jaxpr, consts, *args)
 
     @transform_prim.def_abstract_eval
     def _(*_, inner_jaxpr, **__):
diff --git a/pennylane/workflow/_capture_qnode.py b/pennylane/workflow/_capture_qnode.py
index 6b4c80213e2..a9ed564e63b 100644
--- a/pennylane/workflow/_capture_qnode.py
+++ b/pennylane/workflow/_capture_qnode.py
@@ -215,6 +215,23 @@ def _(*args, qnode, shots, device, execution_config, qfunc_jaxpr, n_consts, batc
     else:
         temp_consts = consts
         temp_args = non_const_args
+
+    # Expand user transforms applied to the qfunc
+    if getattr(qfunc_jaxpr.eqns[0].primitive, "prim_type", "") == "transform":
+        transformed_func = qml.capture.expand_plxpr_transforms(
+            partial(qml.capture.eval_jaxpr, qfunc_jaxpr, temp_consts)
+        )
+
+        qfunc_jaxpr = jax.make_jaxpr(transformed_func)(*temp_args)
+        temp_consts = qfunc_jaxpr.consts
+        qfunc_jaxpr = qfunc_jaxpr.jaxpr
+
+    # Expand user transforms applied to the qnode
+    qfunc_jaxpr = qnode.transform_program(qfunc_jaxpr, temp_consts, *temp_args)
+    temp_consts = qfunc_jaxpr.consts
+    qfunc_jaxpr = qfunc_jaxpr.jaxpr
+
+    # Apply device preprocessing transforms
     qfunc_jaxpr = device_program(qfunc_jaxpr, temp_consts, *temp_args)
     consts = qfunc_jaxpr.consts
     qfunc_jaxpr = qfunc_jaxpr.jaxpr
diff --git a/tests/capture/test_capture_transforms.py b/tests/capture/test_capture_transforms.py
index 15f958ca5c9..5b0a114aa85 100644
--- a/tests/capture/test_capture_transforms.py
+++ b/tests/capture/test_capture_transforms.py
@@ -24,9 +24,9 @@
 
 # pylint: disable=wrong-import-position
 import pennylane as qml
-from pennylane.capture.primitives import cond_prim, for_loop_prim, while_loop_prim
+from pennylane.capture.primitives import cond_prim, for_loop_prim, qnode_prim, while_loop_prim
 from pennylane.tape.plxpr_conversion import CollectOpsandMeas
-from pennylane.transforms.core import TransformError, transform
+from pennylane.transforms.core import TransformError, TransformProgram, transform
 
 pytestmark = [pytest.mark.jax, pytest.mark.usefixtures("enable_disable_plxpr")]
 
@@ -148,31 +148,38 @@ def func(x):
         jaxpr = jax.make_jaxpr(transformed_func)(*args)
         assert (transform_eqn := jaxpr.eqns[0]).primitive == z_to_hadamard._primitive
 
-        params = transform_eqn.params
-        qnode_jaxpr = params["inner_jaxpr"]
-        assert qnode_jaxpr.eqns[0].primitive == qml.capture.qnode_prim
+        qnode_jaxpr = transform_eqn.params["inner_jaxpr"]
+        assert qnode_jaxpr.eqns[0].primitive == qnode_prim
+
+        qnode = qnode_jaxpr.eqns[0].params["qnode"]
+        expected_program = TransformProgram()
+        expected_program.add_transform(z_to_hadamard, *targs, **tkwargs)
+        # Manually change targs from tuple to list
+        expected_program[0]._args = targs  # pylint: disable=protected-access
+        assert qnode.transform_program == expected_program
 
         qfunc_jaxpr = qnode_jaxpr.eqns[0].params["qfunc_jaxpr"]
         expected_jaxpr = jax.make_jaxpr(func)(*args).eqns[0].params["qfunc_jaxpr"]
         for eqn1, eqn2 in zip(qfunc_jaxpr.eqns, expected_jaxpr.eqns, strict=True):
             assert eqn1.primitive == eqn2.primitive
 
-    def test_transform_primitive_eval_not_implemented(self):
-        """Test JAXPR containing a transform primitive cannot be evaluated due to a NotImplementedError."""
+    def test_transform_primitive_eval(self):
+        """Test that evaluating jaxpr containing a transform primitive does not apply
+        the transform."""
 
-        def func(x):
-            y = x * 5
-            return y**0.5
+        def func():
+            return qml.expval(qml.Z(0))
 
-        args = (1.5,)
-        targs = (0, 1)
-        tkwargs = {"dummy_kwarg1": "foo", "dummy_kwarg2": "bar"}
+        targs = (2, 3)
+        tkwargs = {"dummy_kwarg1": "hello", "dummy_kwarg2": "world"}
 
-        transformed_func = z_to_hadamard(func, *targs, **tkwargs)
-        jaxpr = jax.make_jaxpr(transformed_func)(*args)
+        transformed_func = expval_z_obs_to_x_obs(func, *targs, **tkwargs)
+        jaxpr = jax.make_jaxpr(transformed_func)()
 
-        with pytest.raises(NotImplementedError):
-            _ = jax.core.eval_jaxpr(jaxpr.jaxpr, jaxpr.consts, *args)[0]
+        res = jax.core.eval_jaxpr(jaxpr.jaxpr, jaxpr.consts)
+        # If the inner jaxpr was transformed, the output would be <X>, but it should
+        # stay as <Z> if it wasn't transformed
+        assert res == [qml.expval(qml.Z(0))]
 
     def test_multiple_transforms_capture(self):
         """Test that JAXPR containing a transformed qnode primitive is evaluated correctly."""
@@ -232,7 +239,7 @@ def g(i):
             return qml.expval(qml.Z(0))
 
         jaxpr = jax.make_jaxpr(f)()
-        assert jaxpr.eqns[0].primitive == qml.capture.qnode_prim
+        assert jaxpr.eqns[0].primitive == qnode_prim
 
         qfunc_jaxpr = jaxpr.eqns[0].params["qfunc_jaxpr"]
         assert qfunc_jaxpr.eqns[0].primitive == z_to_hadamard._primitive
diff --git a/tests/capture/workflow/test_capture_qnode.py b/tests/capture/workflow/test_capture_qnode.py
index 3273044019d..5efee6e5ee3 100644
--- a/tests/capture/workflow/test_capture_qnode.py
+++ b/tests/capture/workflow/test_capture_qnode.py
@@ -30,6 +30,9 @@
 
 # must be below jax importorskip
 from pennylane.capture.primitives import qnode_prim  # pylint: disable=wrong-import-position
+from pennylane.tape.plxpr_conversion import (  # pylint: disable=wrong-import-position
+    CollectOpsandMeas,
+)
 
 
 def get_qnode_output_eqns(jaxpr):
@@ -351,6 +354,148 @@ def circuit(x):
     assert list(out.keys()) == ["a", "b"]
 
 
+class TestUserTransforms:
+    """Integration tests for applying user transforms to a qnode with program capture."""
+
+    @pytest.mark.unit
+    def test_captured_program_qnode_transform(self):
+        """Test that a transformed qnode is captured correctly."""
+
+        dev = qml.device("default.qubit", wires=3)
+
+        @qml.transforms.cancel_inverses
+        @qml.qnode(dev)
+        def circuit(x):
+            qml.RX(x, 0)
+            qml.X(0)
+            qml.X(0)
+            return qml.expval(qml.Z(0))
+
+        jaxpr = jax.make_jaxpr(circuit)(1.5)
+        # pylint: disable=protected-access
+        assert jaxpr.eqns[0].primitive == qml.transforms.cancel_inverses._primitive
+        inner_jaxpr = jaxpr.eqns[0].params["inner_jaxpr"]
+        assert inner_jaxpr.eqns[0].primitive == qnode_prim
+        qfunc_jaxpr = inner_jaxpr.eqns[0].params["qfunc_jaxpr"]
+
+        collector = CollectOpsandMeas()
+        collector.eval(qfunc_jaxpr, [], 1.5)
+        assert collector.state["ops"] == [qml.RX(1.5, 0), qml.X(0), qml.X(0)]
+        assert collector.state["measurements"] == [qml.expval(qml.Z(0))]
+
+    @pytest.mark.unit
+    def test_captured_program_qfunc_transform(self):
+        """Test that a qnode with a transformed qfunc is captured correctly."""
+
+        dev = qml.device("default.qubit", wires=3)
+
+        @qml.qnode(dev)
+        @qml.transforms.cancel_inverses
+        def circuit(x):
+            qml.RX(x, 0)
+            qml.X(0)
+            qml.X(0)
+            return qml.expval(qml.Z(0))
+
+        jaxpr = jax.make_jaxpr(circuit)(1.5)
+        assert jaxpr.eqns[0].primitive == qnode_prim
+        qfunc_jaxpr = jaxpr.eqns[0].params["qfunc_jaxpr"]
+        # pylint: disable=protected-access
+        assert qfunc_jaxpr.eqns[0].primitive == qml.transforms.cancel_inverses._primitive
+
+        inner_jaxpr = qfunc_jaxpr.eqns[0].params["inner_jaxpr"]
+        collector = CollectOpsandMeas()
+        collector.eval(inner_jaxpr, [], 1.5)
+        assert collector.state["ops"] == [qml.RX(1.5, 0), qml.X(0), qml.X(0)]
+        assert collector.state["measurements"] == [qml.expval(qml.Z(0))]
+
+    @pytest.mark.unit
+    def test_captured_program_qnode_qfunc_transform(self):
+        """Test that a transformed qnode with a transformed qfunc is captured correctly."""
+
+        dev = qml.device("default.qubit", wires=3)
+
+        @qml.transforms.cancel_inverses
+        @qml.qnode(dev)
+        @qml.transforms.merge_rotations
+        def circuit(x):
+            qml.RX(x, 0)
+            qml.X(0)
+            qml.X(0)
+            return qml.expval(qml.Z(0))
+
+        jaxpr = jax.make_jaxpr(circuit)(1.5)
+        # pylint: disable=protected-access
+        assert jaxpr.eqns[0].primitive == qml.transforms.cancel_inverses._primitive
+        inner_jaxpr = jaxpr.eqns[0].params["inner_jaxpr"]
+        assert inner_jaxpr.eqns[0].primitive == qnode_prim
+        qfunc_jaxpr = inner_jaxpr.eqns[0].params["qfunc_jaxpr"]
+
+        assert qfunc_jaxpr.eqns[0].primitive == qml.transforms.merge_rotations._primitive
+        inner_jaxpr2 = qfunc_jaxpr.eqns[0].params["inner_jaxpr"]
+
+        collector = CollectOpsandMeas()
+        collector.eval(inner_jaxpr2, [], 1.5)
+        assert collector.state["ops"] == [qml.RX(1.5, 0), qml.X(0), qml.X(0)]
+        assert collector.state["measurements"] == [qml.expval(qml.Z(0))]
+
+    @pytest.mark.unit
+    def test_device_jaxpr(self, monkeypatch):
+        """Test that jaxpr recieved by a device when executing a transformed qnode has been
+        transformed appropriately."""
+
+        device_jaxpr = None
+
+        def dummy_eval_jaxpr(
+            jaxpr, consts, *args, execution_config
+        ):  # pylint: disable=unused-argument
+            nonlocal device_jaxpr
+            device_jaxpr = jaxpr
+            return [1.0]
+
+        dev = qml.device("default.qubit", wires=3)
+        monkeypatch.setattr(dev, "eval_jaxpr", dummy_eval_jaxpr)
+
+        @partial(qml.transforms.decompose, gate_set=[qml.RX, qml.RY, qml.RZ])
+        @qml.qnode(dev)
+        @qml.transforms.cancel_inverses
+        def circuit(x, y, z):
+            qml.Rot(x, y, z, 0)
+            qml.X(0)
+            qml.X(0)
+            return qml.expval(qml.Z(0))
+
+        _ = circuit(1.5, 2.5, 3.5)
+        assert all(
+            getattr(eqn.primitive, "prim_type", "") != "transform" for eqn in device_jaxpr.eqns
+        )
+        assert device_jaxpr.eqns[0].primitive == qml.RZ._primitive
+        assert device_jaxpr.eqns[1].primitive == qml.RY._primitive
+        assert device_jaxpr.eqns[2].primitive == qml.RZ._primitive
+        assert device_jaxpr.eqns[3].primitive == qml.PauliZ._primitive
+        assert device_jaxpr.eqns[4].primitive == qml.measurements.ExpectationMP._obs_primitive
+
+    @pytest.mark.integration
+    def test_execution(self):
+        """Test that a transformed qnode is executed correctly."""
+
+        dev = qml.device("default.qubit", wires=3)
+
+        @qml.transforms.cancel_inverses
+        @qml.qnode(dev)
+        @qml.transforms.merge_rotations
+        def circuit(x):
+            qml.RX(x, 0)
+            qml.RX(4 * x, 0)
+            qml.X(0)
+            qml.X(0)
+            return qml.expval(qml.Z(0))
+
+        res = circuit(1.5)
+        expected = jnp.cos(5 * 1.5)
+        assert jnp.allclose(res, expected)
+
+
 @pytest.mark.parametrize("dev_name", ["default.qubit", "lightning.qubit"])
 class TestDevicePreprocessing:
     """Integration tests for preprocessing and executing qnodes with program capture."""


Reviews:
user: login: andrijapau
body: Looks good so far, just an initial comment.
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-04-03T20:16:43Z
user: login: astralcai
body: Looks straightforward enough, no questions from me.
state: APPROVED
author_association: CONTRIBUTOR
submitted_at: 2025-04-04T13:55:53Z
user: login: mudit2812
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-04-04T14:11:22Z
user: login: andrijapau
body: Looks good to me. Just some non-blocking comments! 🚀 
state: APPROVED
author_association: CONTRIBUTOR
submitted_at: 2025-04-04T14:18:35Z
user: login: mudit2812
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-04-04T14:28:17Z
user: login: mudit2812
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-04-04T14:28:47Z
user: login: mudit2812
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-04-04T14:29:56Z
user: login: mudit2812
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-04-04T18:35:19Z
user: login: mudit2812
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-04-04T18:40:52Z
user: login: andrijapau
body: Could use some clarification here 😄 
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-04-04T19:01:22Z
user: login: astralcai
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-04-04T19:01:22Z
user: login: mudit2812
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-04-04T19:07:59Z
user: login: mudit2812
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-04-04T19:09:31Z
user: login: mudit2812
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-04-04T19:10:40Z
user: login: astralcai
state: APPROVED
author_association: CONTRIBUTOR
submitted_at: 2025-04-04T19:12:23Z
user: login: andrijapau
body: Looks good to me! 🚀  Thanks for the quick change.
state: APPROVED
author_association: CONTRIBUTOR
submitted_at: 2025-04-04T19:20:14Z

Review comments:
diff_hunk: @@ -319,6 +319,12 @@ def default_qnode_transform(self, qnode, targs, tkwargs):
         )
         return qnode
 
+    def _capture_qnode_transform(self, qnode, targs, tkwargs):
+        """Apply the transform to the quantum function of a qnode when capture is enabled."""
+
+        new_qfunc = self._capture_callable_transform(qnode.func, targs, tkwargs)
+        return qnode.update(func=new_qfunc)
path: pennylane/transforms/core/transform_dispatcher.py
user: login: andrijapau
body: Love seeing update in the wild ❤️ 
created_at: 2025-04-03T20:13:15Z
author_association: CONTRIBUTOR
diff_hunk: @@ -319,6 +319,12 @@ def default_qnode_transform(self, qnode, targs, tkwargs):
         )
         return qnode
 
+    def _capture_qnode_transform(self, qnode, targs, tkwargs):
+        """Apply the transform to the quantum function of a qnode when capture is enabled."""
+
+        new_qfunc = self._capture_callable_transform(qnode.func, targs, tkwargs)
+        return qnode.update(func=new_qfunc)
path: pennylane/transforms/core/transform_dispatcher.py
user: login: andrijapau
body: But do we really need this new private method? Why don't we just put these two lines after L193?
created_at: 2025-04-03T20:14:17Z
author_association: CONTRIBUTOR
diff_hunk: @@ -161,6 +161,10 @@ With `qml.decompositions.enable_graph()`, the following new features are availab
 
 <h4>Capturing and Representing Hybrid Programs 📥</h4>
 
+* Applying transforms to a `QNode` now correctly transforms it with program capture enabled, rather than raising a
+  `NotImplementedError`.
path: doc/releases/changelog-dev.md
user: login: andrijapau
body: 🤔 
```suggestion
* Transformations can now be directly applied to a `QNode` with program capture enabled.
```
created_at: 2025-04-04T14:09:16Z
author_association: CONTRIBUTOR
diff_hunk: @@ -319,6 +319,12 @@ def default_qnode_transform(self, qnode, targs, tkwargs):
         )
         return qnode
 
+    def _capture_qnode_transform(self, qnode, targs, tkwargs):
+        """Apply the transform to the quantum function of a qnode when capture is enabled."""
+
+        new_qfunc = self._capture_callable_transform(qnode.func, targs, tkwargs)
+        return qnode.update(func=new_qfunc)
path: pennylane/transforms/core/transform_dispatcher.py
user: login: mudit2812
body: Fixed.
created_at: 2025-04-04T14:11:22Z
author_association: CONTRIBUTOR
diff_hunk: @@ -191,7 +191,8 @@ def processing_fn(results):
 
         if isinstance(obj, qml.QNode):
             if qml.capture.enabled():
-                return self._capture_callable_transform(obj, targs, tkwargs)
+                new_qfunc = self._capture_callable_transform(obj.func, targs, tkwargs)
+                return obj.update(func=new_qfunc)
path: pennylane/transforms/core/transform_dispatcher.py
user: login: andrijapau
body: This will create a new `QNode` object, are we okay with that?
created_at: 2025-04-04T14:12:16Z
author_association: CONTRIBUTOR
diff_hunk: @@ -351,6 +354,148 @@ def circuit(x):
     assert list(out.keys()) == ["a", "b"]
 
 
+class TestUserTransforms:
+    """Integration tests for applying user transforms to a qnode with program capture."""
+
+    @pytest.mark.unit
+    def test_captured_program_qnode_transform(self):
+        """Test that a transformed qnode is captured correctly."""
+
+        dev = qml.device("default.qubit", wires=3)
+
+        @qml.transforms.cancel_inverses
+        @qml.qnode(dev)
+        def circuit(x):
+            qml.RX(x, 0)
+            qml.X(0)
+            qml.X(0)
+            return qml.expval(qml.Z(0))
+
+        jaxpr = jax.make_jaxpr(circuit)(1.5)
+        assert jaxpr.eqns[0].primitive == qnode_prim
+        qfunc_jaxpr = jaxpr.eqns[0].params["qfunc_jaxpr"]
+        # pylint: disable=protected-access
+        assert qfunc_jaxpr.eqns[0].primitive == qml.transforms.cancel_inverses._primitive
+
+        inner_jaxpr = qfunc_jaxpr.eqns[0].params["inner_jaxpr"]
+        collector = CollectOpsandMeas()
+        collector.eval(inner_jaxpr, [], 1.5)
+        assert collector.state["ops"] == [qml.RX(1.5, 0), qml.X(0), qml.X(0)]
+        assert collector.state["measurements"] == [qml.expval(qml.Z(0))]
+
+    @pytest.mark.unit
+    def test_captured_program_qfunc_transform(self):
+        """Test that a qnode with a transformed qfunc is captured correctly."""
+
+        dev = qml.device("default.qubit", wires=3)
+
+        @qml.qnode(dev)
+        @qml.transforms.cancel_inverses
+        def circuit(x):
+            qml.RX(x, 0)
+            qml.X(0)
+            qml.X(0)
+            return qml.expval(qml.Z(0))
+
+        jaxpr = jax.make_jaxpr(circuit)(1.5)
+        assert jaxpr.eqns[0].primitive == qnode_prim
+        qfunc_jaxpr = jaxpr.eqns[0].params["qfunc_jaxpr"]
+        # pylint: disable=protected-access
+        assert qfunc_jaxpr.eqns[0].primitive == qml.transforms.cancel_inverses._primitive
+
+        inner_jaxpr = qfunc_jaxpr.eqns[0].params["inner_jaxpr"]
+        collector = CollectOpsandMeas()
+        collector.eval(inner_jaxpr, [], 1.5)
+        assert collector.state["ops"] == [qml.RX(1.5, 0), qml.X(0), qml.X(0)]
+        assert collector.state["measurements"] == [qml.expval(qml.Z(0))]
+
+    @pytest.mark.unit
+    def test_captured_program_qnode_qfunc_transform(self):
+        """Test that a transformed qnode with a transformed qfunc is captured correctly."""
+
+        dev = qml.device("default.qubit", wires=3)
+
+        @qml.transforms.cancel_inverses
+        @qml.qnode(dev)
+        @qml.transforms.merge_rotations
+        def circuit(x):
+            qml.RX(x, 0)
+            qml.X(0)
+            qml.X(0)
+            return qml.expval(qml.Z(0))
+
+        jaxpr = jax.make_jaxpr(circuit)(1.5)
+        assert jaxpr.eqns[0].primitive == qnode_prim
+        qfunc_jaxpr = jaxpr.eqns[0].params["qfunc_jaxpr"]
+        # pylint: disable=protected-access
+        assert qfunc_jaxpr.eqns[0].primitive == qml.transforms.cancel_inverses._primitive
+        inner_jaxpr = qfunc_jaxpr.eqns[0].params["inner_jaxpr"]
+        assert inner_jaxpr.eqns[0].primitive == qml.transforms.merge_rotations._primitive
+        inner_jaxpr2 = inner_jaxpr.eqns[0].params["inner_jaxpr"]
+
+        collector = CollectOpsandMeas()
+        collector.eval(inner_jaxpr2, [], 1.5)
+        assert collector.state["ops"] == [qml.RX(1.5, 0), qml.X(0), qml.X(0)]
+        assert collector.state["measurements"] == [qml.expval(qml.Z(0))]
+
+    @pytest.mark.unit
+    def test_device_jaxpr(self, monkeypatch):
+        """Test that jaxpr recieved by a device when executing a transformed qnode has been
+        transformed appropriately."""
+
+        device_jaxpr = None
+
+        def dummy_eval_jaxpr(
+            jaxpr, consts, *args, execution_config
+        ):  # pylint: disable=unused-argument
+            nonlocal device_jaxpr
+            device_jaxpr = jaxpr
+            return [1.0]
+
+        dev = qml.device("default.qubit", wires=3)
+        monkeypatch.setattr(dev, "eval_jaxpr", dummy_eval_jaxpr)
+
+        @partial(qml.transforms.decompose, gate_set=[qml.RX, qml.RY, qml.RZ])
+        @qml.qnode(dev)
+        @qml.transforms.cancel_inverses
+        def circuit(x, y, z):
+            qml.Rot(x, y, z, 0)
+            qml.X(0)
+            qml.X(0)
+            return qml.expval(qml.Z(0))
+
+        _ = circuit(1.5, 2.5, 3.5)
+        print(device_jaxpr)
path: tests/capture/workflow/test_capture_qnode.py
user: login: andrijapau
body: ```suggestion
```
created_at: 2025-04-04T14:15:13Z
author_association: CONTRIBUTOR
path: tests/capture/workflow/test_capture_qnode.py
user: login: andrijapau
body: Do you think it's worth it to add a test that applies the following: (transform, qnode, transform, transform). The difference being that there are two user transforms?
created_at: 2025-04-04T14:18:32Z
author_association: CONTRIBUTOR
diff_hunk: @@ -161,6 +161,10 @@ With `qml.decompositions.enable_graph()`, the following new features are availab
 
 <h4>Capturing and Representing Hybrid Programs 📥</h4>
 
+* Applying transforms to a `QNode` now correctly transforms it with program capture enabled, rather than raising a
+  `NotImplementedError`.
path: doc/releases/changelog-dev.md
user: login: mudit2812
body: How about adding "without the use of the `@qml.capture.expand_plxpr_transforms` decorator" to what you suggested?
created_at: 2025-04-04T14:28:17Z
author_association: CONTRIBUTOR
diff_hunk: @@ -191,7 +191,8 @@ def processing_fn(results):
 
         if isinstance(obj, qml.QNode):
             if qml.capture.enabled():
-                return self._capture_callable_transform(obj, targs, tkwargs)
+                new_qfunc = self._capture_callable_transform(obj.func, targs, tkwargs)
+                return obj.update(func=new_qfunc)
path: pennylane/transforms/core/transform_dispatcher.py
user: login: mudit2812
body: That is consistent with the behaviour when capture is disabled. I think it's a good idea for us to generally move away from mutating qnodes.
created_at: 2025-04-04T14:28:47Z
author_association: CONTRIBUTOR
path: tests/capture/workflow/test_capture_qnode.py
user: login: mudit2812
body: Sure, better safe than sorry. I'll also add a test for (transform, transform, qnode, transform).
created_at: 2025-04-04T14:29:56Z
author_association: CONTRIBUTOR
diff_hunk: @@ -585,7 +586,10 @@ def _create_transform_primitive(name):
     def _(
         *all_args, inner_jaxpr, args_slice, consts_slice, targs_slice, tkwargs
     ):  # pylint: disable=unused-argument
-        raise NotImplementedError
+        args = all_args[args_slice]
+        consts = all_args[consts_slice]
+
+        return qml.capture.eval_jaxpr(inner_jaxpr, consts, *args)
path: pennylane/transforms/core/transform_dispatcher.py
user: login: mudit2812
body: This was needed so that calling a transformed function does not raise a `NotImplementedError`.
created_at: 2025-04-04T18:35:19Z
author_association: CONTRIBUTOR
path: tests/capture/workflow/test_capture_qnode.py
user: login: mudit2812
body: With the change in implementation, adding these tests doesn't make sense to me anymore.
created_at: 2025-04-04T18:40:52Z
author_association: CONTRIBUTOR
diff_hunk: @@ -215,6 +215,23 @@ def _(*args, qnode, shots, device, execution_config, qfunc_jaxpr, n_consts, batc
     else:
         temp_consts = consts
         temp_args = non_const_args
+
+    # Expand user transforms applied to the qfunc
+    if getattr(qfunc_jaxpr.eqns[0].primitive, "prim_type", "") == "transform":
+        transformed_func = qml.capture.expand_plxpr_transforms(
+            partial(qml.capture.eval_jaxpr, qfunc_jaxpr, temp_consts)
+        )
+
+        qfunc_jaxpr = jax.make_jaxpr(transformed_func)(*temp_args)
+        temp_consts = qfunc_jaxpr.consts
+        qfunc_jaxpr = qfunc_jaxpr.jaxpr
+
+    # Expand user transforms applied to the qnode
+    qfunc_jaxpr = qnode.transform_program(qfunc_jaxpr, temp_consts, *temp_args)
+    temp_consts = qfunc_jaxpr.consts
+    qfunc_jaxpr = qfunc_jaxpr.jaxpr
+
path: pennylane/workflow/_capture_qnode.py
user: login: andrijapau
body: Remind me how this expands the user transforms applied to the qnode?
created_at: 2025-04-04T18:57:51Z
author_association: CONTRIBUTOR
diff_hunk: @@ -148,31 +148,38 @@ def func(x):
         jaxpr = jax.make_jaxpr(transformed_func)(*args)
         assert (transform_eqn := jaxpr.eqns[0]).primitive == z_to_hadamard._primitive
 
-        params = transform_eqn.params
-        qnode_jaxpr = params["inner_jaxpr"]
-        assert qnode_jaxpr.eqns[0].primitive == qml.capture.qnode_prim
+        qnode_jaxpr = transform_eqn.params["inner_jaxpr"]
+        assert qnode_jaxpr.eqns[0].primitive == qnode_prim
+
+        qnode = qnode_jaxpr.eqns[0].params["qnode"]
+        expected_program = TransformProgram()
+        expected_program.add_transform(z_to_hadamard, *targs, **tkwargs)
+        # Manually change targs from tuple to list
+        expected_program[0]._args = targs  # pylint: disable=protected-access
+        assert qnode.transform_program == expected_program
path: tests/capture/test_capture_transforms.py
user: login: andrijapau
body: 🤔 Is this a problem?
created_at: 2025-04-04T19:00:03Z
author_association: CONTRIBUTOR
diff_hunk: @@ -191,7 +191,8 @@ def processing_fn(results):
 
         if isinstance(obj, qml.QNode):
             if qml.capture.enabled():
-                return self._capture_callable_transform(obj, targs, tkwargs)
+                new_qnode = self.default_qnode_transform(obj, targs, tkwargs)
+                return self._capture_callable_transform(new_qnode, targs, tkwargs)
path: pennylane/transforms/core/transform_dispatcher.py
user: login: astralcai
body: I may not understand this correctly but shouldn't we return a qnode here?
created_at: 2025-04-04T19:01:22Z
author_association: CONTRIBUTOR
diff_hunk: @@ -148,31 +148,38 @@ def func(x):
         jaxpr = jax.make_jaxpr(transformed_func)(*args)
         assert (transform_eqn := jaxpr.eqns[0]).primitive == z_to_hadamard._primitive
 
-        params = transform_eqn.params
-        qnode_jaxpr = params["inner_jaxpr"]
-        assert qnode_jaxpr.eqns[0].primitive == qml.capture.qnode_prim
+        qnode_jaxpr = transform_eqn.params["inner_jaxpr"]
+        assert qnode_jaxpr.eqns[0].primitive == qnode_prim
+
+        qnode = qnode_jaxpr.eqns[0].params["qnode"]
+        expected_program = TransformProgram()
+        expected_program.add_transform(z_to_hadamard, *targs, **tkwargs)
+        # Manually change targs from tuple to list
+        expected_program[0]._args = targs  # pylint: disable=protected-access
+        assert qnode.transform_program == expected_program
path: tests/capture/test_capture_transforms.py
user: login: mudit2812
body: No, it's just because of how `targs` are handled in `TransformDispatcher.__call__`. [This line](https://github.com/PennyLaneAI/pennylane/blob/83a585b728374ff625c0d614741c07bc2af692fa/pennylane/transforms/core/transform_dispatcher.py#L164) changes it from a tuple to a list.
created_at: 2025-04-04T19:07:59Z
author_association: CONTRIBUTOR
diff_hunk: @@ -215,6 +215,23 @@ def _(*args, qnode, shots, device, execution_config, qfunc_jaxpr, n_consts, batc
     else:
         temp_consts = consts
         temp_args = non_const_args
+
+    # Expand user transforms applied to the qfunc
+    if getattr(qfunc_jaxpr.eqns[0].primitive, "prim_type", "") == "transform":
+        transformed_func = qml.capture.expand_plxpr_transforms(
+            partial(qml.capture.eval_jaxpr, qfunc_jaxpr, temp_consts)
+        )
+
+        qfunc_jaxpr = jax.make_jaxpr(transformed_func)(*temp_args)
+        temp_consts = qfunc_jaxpr.consts
+        qfunc_jaxpr = qfunc_jaxpr.jaxpr
+
+    # Expand user transforms applied to the qnode
+    qfunc_jaxpr = qnode.transform_program(qfunc_jaxpr, temp_consts, *temp_args)
+    temp_consts = qfunc_jaxpr.consts
+    qfunc_jaxpr = qfunc_jaxpr.jaxpr
+
path: pennylane/workflow/_capture_qnode.py
user: login: mudit2812
body: User transforms applied to the qnode will be in the qnode's transform program.
created_at: 2025-04-04T19:09:31Z
author_association: CONTRIBUTOR
diff_hunk: @@ -191,7 +191,8 @@ def processing_fn(results):
 
         if isinstance(obj, qml.QNode):
             if qml.capture.enabled():
-                return self._capture_callable_transform(obj, targs, tkwargs)
+                new_qnode = self.default_qnode_transform(obj, targs, tkwargs)
+                return self._capture_callable_transform(new_qnode, targs, tkwargs)
path: pennylane/transforms/core/transform_dispatcher.py
user: login: mudit2812
body: With capture enabled, a callable is always returned. This was the case before this PR as well. If we return a qnode, then the first primitive in the captured program of a qnode will be the qnode primitive. But, when a transform is applied to a qnode, the first primitive should be the transform's primitive.
created_at: 2025-04-04T19:10:40Z
author_association: CONTRIBUTOR

PR comments:
user: login: codecov[bot]
created_at: 2025-04-04T07:54:17Z
body: ## [Codecov](https://app.codecov.io/gh/PennyLaneAI/pennylane/pull/7199?dropdown=coverage&src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) Report
All modified and coverable lines are covered by tests :white_check_mark:
> Project coverage is 99.64%. Comparing base [(`85f6ff7`)](https://app.codecov.io/gh/PennyLaneAI/pennylane/commit/85f6ff7f65629f850e086d6a5189464cf04845b8?dropdown=coverage&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) to head [(`eb29086`)](https://app.codecov.io/gh/PennyLaneAI/pennylane/commit/eb29086906dd4693ffb888bd31d5c8a6cfe4ad37?dropdown=coverage&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI).
> Report is 342 commits behind head on master.

<details><summary>Additional details and impacted files</summary>


```diff
@@           Coverage Diff           @@
##           master    #7199   +/-   ##
=======================================
  Coverage   99.64%   99.64%           
=======================================
  Files         511      511           
  Lines       49885    49896   +11     
=======================================
+ Hits        49709    49720   +11     
  Misses        176      176           
```

</details>

[:umbrella: View full report in Codecov by Sentry](https://app.codecov.io/gh/PennyLaneAI/pennylane/pull/7199?dropdown=coverage&src=pr&el=continue&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI).   
:loudspeaker: Have feedback on the report? [Share it here](https://about.codecov.io/codecov-pr-comment-feedback/?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI).

<details><summary> :rocket: New features to boost your workflow: </summary>

- :snowflake: [Test Analytics](https://docs.codecov.com/docs/test-analytics): Detect flaky tests, report on failures, and find test suite problems.
</details>
author_association: NONE
user: login: mudit2812
created_at: 2025-04-04T14:04:35Z
body: @isaacdevlugt if this PR is merged, we will no longer have to manually use `qml.capture.expand_plxpr_transforms`, and that info can be removed from the sharp bits page.
author_association: CONTRIBUTOR