[Capture] add eval_jaxpr method to null.qubit

doc/releases/changelog-dev.md

@@ -154,6 +154,9 @@
 * The `qml.clifford_t_decomposition` has been improved to use less gates when decomposing `qml.PhaseShift`.
   [(#6842)](https://github.com/PennyLaneAI/pennylane/pull/6842)
 
+* `null.qubit` can now execute jaxpr.
+  [(#6924)](https://github.com/PennyLaneAI/pennylane/pull/6924)
+
 <h3>Breaking changes 💔</h3>
 
 * `MultiControlledX` no longer accepts strings as control values.

pennylane/devices/null_qubit.py

@@ -404,3 +404,17 @@ def execute_and_compute_vjp(
         results = tuple(self._simulate(c, _interface(execution_config)) for c in circuits)
         vjps = tuple(self._vjp(c, _interface(execution_config)) for c in circuits)
         return results, vjps
+
+    def eval_jaxpr(self, jaxpr: "jax.core.Jaxpr", consts: list, *args) -> list:
+        from pennylane.capture.primitives import (  # pylint: disable=import-outside-toplevel
+            AbstractMeasurement,
+        )
+
+        def zeros_like(var):
+            if isinstance(var.aval, AbstractMeasurement):
+                shots = self.shots.total_shots
+                s, dtype = var.aval.abstract_eval(num_device_wires=len(self.wires), shots=shots)
+                return math.zeros(s, dtype=dtype, like="jax")
+            return math.zeros(var.aval.shape, dtype=var.aval.dtype, like="jax")
+
+        return [zeros_like(var) for var in jaxpr.outvars]

tests/devices/test_null_qubit.py

@@ -1273,3 +1273,42 @@ def circuit(param):
     res = qml.QNode(circuit, nq)(x)
     target = qml.QNode(circuit, dq)(x)
     assert qml.math.shape(res) == qml.math.shape(target)
+
+
+# pylint: disable=unused-argument
+@pytest.mark.jax
+def test_execute_plxpr(enable_disable_plxpr):
+    """Test that null.qubit can execute plxpr."""
+
+    import jax
+
+    def f(x):
+        qml.RX(x, 0)
+        return qml.expval(qml.Z(0)), qml.probs(), 4, qml.var(qml.X(0)), qml.state()
+
+    jaxpr = jax.make_jaxpr(f)(0.5)
+
+    dev = qml.device("null.qubit", wires=4)
+    res = dev.eval_jaxpr(jaxpr.jaxpr, jaxpr.consts, 1.0)
+    assert qml.math.allclose(res[0], 0)
+    assert qml.math.allclose(res[1], jax.numpy.zeros(2**4))
+    assert qml.math.allclose(res[2], 0)  # other values are still just zero
+    assert qml.math.allclose(res[3], 0)
+    assert qml.math.allclose(res[4], jax.numpy.zeros(2**4, dtype=complex))
+
+
+@pytest.mark.jax
+def test_execute_plxpr_shots(enable_disable_plxpr):
+    import jax
+
+    def f(x):
+        qml.RX(x, 0)
+        return qml.expval(qml.Z(0)), 5, qml.sample(wires=(0, 1))
+
+    jaxpr = jax.make_jaxpr(f)(0.5)
+
+    dev = qml.device("null.qubit", wires=4, shots=50)
+    res = dev.eval_jaxpr(jaxpr.jaxpr, jaxpr.consts, 1.0)
+    assert qml.math.allclose(res[0], 0)
+    assert qml.math.allclose(res[1], 0)
+    assert qml.math.allclose(res[2], jax.numpy.zeros((50, 2)))