[Capture] add eval_jaxpr method to DefaultQubit

doc/releases/changelog-dev.md

@@ -49,7 +49,9 @@
   pennylane variant jaxpr.
   [(#6141)](https://github.com/PennyLaneAI/pennylane/pull/6141)
 
-* A `DefaultQubitInterpreter` class has been added to provide plxpr execution using python based tools.
+* A `DefaultQubitInterpreter` class has been added to provide plxpr execution using python based tools,
+  and the `DefaultQubit.eval_jaxpr` method is now implemented.
+  [(#6594)](https://github.com/PennyLaneAI/pennylane/pull/6594)
   [(#6328)](https://github.com/PennyLaneAI/pennylane/pull/6328)
 
 * An optional method `eval_jaxpr` is added to the device API for native execution of plxpr programs.

pennylane/devices/default_qubit.py

@@ -32,7 +32,7 @@
 from pennylane.tape import QuantumScript, QuantumScriptBatch, QuantumScriptOrBatch
 from pennylane.transforms import convert_to_numpy_parameters
 from pennylane.transforms.core import TransformProgram
-from pennylane.typing import PostprocessingFn, Result, ResultBatch
+from pennylane.typing import PostprocessingFn, Result, ResultBatch, TensorLike
 
 from . import Device
 from .execution_config import DefaultExecutionConfig, ExecutionConfig
@@ -891,6 +891,28 @@ def execute_and_compute_vjp(
 
         return tuple(zip(*results))
 
+    # pylint: disable=import-outside-toplevel
+    def eval_jaxpr(
+        self, jaxpr: "jax.core.Jaxpr", consts: list[TensorLike], *args
+    ) -> list[TensorLike]:
+        from .qubit.dq_interpreter import DefaultQubitInterpreter
+
+        if self.wires is None:
+            raise qml.DeviceError("Device wires are required for jaxpr execution.")
+        if self.shots.has_partitioned_shots:
+            raise qml.DeviceError("Shot vectors are unsupported with jaxpr execution.")
+        if self._prng_key is not None:
+            key = self.get_prng_keys()[0]
+        else:
+            import jax
+
+            key = jax.random.PRNGKey(self._rng.integers(100000))
+
+        interpreter = DefaultQubitInterpreter(
+            num_wires=len(self.wires), shots=self.shots.total_shots, key=key
+        )
+        return interpreter.eval(jaxpr, consts, *args)
+
 
 def _simulate_wrapper(circuit, kwargs):
     return simulate(circuit, **kwargs)

tests/devices/default_qubit/test_default_qubit_plxpr.py

@@ -0,0 +1,97 @@
+# Copyright 2024 Xanadu Quantum Technologies Inc.
+
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+
+#     http://www.apache.org/licenses/LICENSE-2.0
+
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""Tests for default qubit executing jaxpr."""
+
+import pytest
+
+import pennylane as qml
+
+jax = pytest.importorskip("jax")
+pytestmark = pytest.mark.jax
+
+
+@pytest.fixture(autouse=True)
+def enable_disable_plxpr():
+    qml.capture.enable()
+    yield
+    qml.capture.disable()
+
+
+def test_requires_wires():
+    """Test that a device error is raised if device wires are not specified."""
+
+    jaxpr = jax.make_jaxpr(lambda x: x + 1)(0.1)
+    dev = qml.device("default.qubit")
+
+    with pytest.raises(qml.DeviceError, match="Device wires are required."):
+        dev.eval_jaxpr(jaxpr.jaxpr, jaxpr.consts, 0.2)
+
+
+def test_no_partitioned_shots():
+    """Test that an error is raised if the device has partitioned shots."""
+
+    jaxpr = jax.make_jaxpr(lambda x: x + 1)(0.1)
+    dev = qml.device("default.qubit", wires=1, shots=(100, 100))
+
+    with pytest.raises(qml.DeviceError, match="Shot vectors are unsupported with jaxpr execution."):
+        dev.eval_jaxpr(jaxpr.jaxpr, jaxpr.consts, 0.2)
+
+
+def test_use_device_prng():
+    """Test that sampling depends on the device prng."""
+
+    key1 = jax.random.PRNGKey(1234)
+    key2 = jax.random.PRNGKey(1234)
+
+    dev1 = qml.device("default.qubit", wires=1, shots=100, seed=key1)
+    dev2 = qml.device("default.qubit", wires=1, shots=100, seed=key2)
+
+    def f():
+        qml.H(0)
+        return qml.sample(wires=0)
+
+    jaxpr = jax.make_jaxpr(f)()
+
+    samples1 = dev1.eval_jaxpr(jaxpr.jaxpr, jaxpr.consts)
+    samples2 = dev2.eval_jaxpr(jaxpr.jaxpr, jaxpr.consts)
+
+    assert qml.math.allclose(samples1, samples2)
+
+
+def test_no_prng_key():
+    """Test that that sampling works without a provided prng key."""
+
+    dev = qml.device("default.qubit", wires=1, shots=100)
+
+    def f():
+        return qml.sample(wires=0)
+
+    jaxpr = jax.make_jaxpr(f)()
+    res = dev.eval_jaxpr(jaxpr.jaxpr, jaxpr.consts)
+    assert qml.math.allclose(res, jax.numpy.zeros(100))
+
+
+def test_simple_execution():
+    """Test the execution, jitting, and gradient of a simple quantum circuit."""
+
+    def f(x):
+        qml.RX(x, 0)
+        return qml.expval(qml.Z(0))
+
+    jaxpr = jax.make_jaxpr(f)(0.123)
+
+    dev = qml.device("default.qubit", wires=1)
+
+    res = dev.eval_jaxpr(jaxpr.jaxpr, jaxpr.consts, 0.5)
+    assert qml.math.allclose(res, jax.numpy.cos(0.5))