[Capture] JIT entire workflows with default qubit

doc/releases/changelog-dev.md

@@ -165,6 +165,7 @@ featuring a `simulate` function for simulating mixed states in analytic mode.
 * Execution with capture enabled now follows a new execution pipeline and natively passes the
   captured jaxpr to the device. Since it no longer falls back to the old pipeline, execution
   only works with a reduced feature set.
+  [(#6496)](https://github.com/PennyLaneAI/pennylane/pull/6596)
   [(#6596)](https://github.com/PennyLaneAI/pennylane/pull/6596)
 
 * PennyLane transforms can now be captured as primitives with experimental program capture enabled.

pennylane/workflow/_capture_qnode.py

@@ -119,7 +119,7 @@
 has_jax = True
 try:
     import jax
-    from jax.interpreters import ad, batching
+    from jax.interpreters import ad, batching, mlir
 except ImportError:
     has_jax = False
 
@@ -298,6 +298,8 @@
 
     batching.primitive_batchers[qnode_prim] = _qnode_batching_rule
 
+    mlir.register_lowering(qnode_prim, mlir.lower_fun(qnode_impl, multiple_results=True))
+
     return qnode_prim
 
 

tests/capture/test_capture_mid_measure.py

@@ -309,7 +309,7 @@ class TestMidMeasureExecute:
 
     @pytest.mark.xfail(strict=False)  # single branch statistics sometimes gives good results
     @pytest.mark.parametrize("reset", [True, False])
-    @pytest.mark.parametrize("postselect", [None, 0, 1])
+    @pytest.mark.parametrize("postselect", [pytest.param(None, marks=pytest.mark.xfail), 0, 1])
     @pytest.mark.parametrize("phi", jnp.arange(1.0, 2 * jnp.pi, 1.5))
     def test_simple_circuit_execution(self, phi, reset, postselect, get_device, shots, mp_fn, seed):
         """Test that circuits with mid-circuit measurements can be executed in a QNode."""

tests/capture/test_capture_qnode.py

@@ -353,6 +353,19 @@ def circuit(x):
     assert qml.math.allclose(jvp, (qml.math.cos(x), -qml.math.sin(x) * xt))
 
 
+def test_qnode_jit():
+    """Test that executions on default qubit can be jitted."""
+
+    @qml.qnode(qml.device("default.qubit", wires=1))
+    def circuit(x):
+        qml.RX(x, 0)
+        return qml.expval(qml.Z(0))
+
+    x = jax.numpy.array(-0.5)
+    res = jax.jit(circuit)(0.5)
+    assert qml.math.allclose(res, jax.numpy.cos(x))
+
+
 # pylint: disable=too-many-public-methods
 class TestQNodeVmapIntegration:
     """Tests for integrating JAX vmap with the QNode primitive."""