Capture a qnode into jaxpr

doc/releases/changelog-dev.md

@@ -93,9 +93,10 @@
   `m = measure(0); qml.sample(m)`.
   [(#5673)](https://github.com/PennyLaneAI/pennylane/pull/5673)
 
-* PennyLane operators can now automatically be captured as instructions in JAXPR. See the experimental
-  `capture` module for more information.
+* PennyLane operators, measurements, and QNodes can now automatically be captured as instructions in JAXPR.
+  [(#5564)](https://github.com/PennyLaneAI/pennylane/pull/5564)
   [(#5511)](https://github.com/PennyLaneAI/pennylane/pull/5511)
+  [(#5708)](https://github.com/PennyLaneAI/pennylane/pull/5708)
 
 * The `decompose` transform has an `error` kwarg to specify the type of error that should be raised, 
   allowing error types to be more consistent with the context the `decompose` function is used in.

pennylane/capture/__init__.py

@@ -21,6 +21,19 @@
 
     This module is experimental and will change significantly in the future.
 
+.. currentmodule:: pennylane.capture
+
+.. autosummary::
+    :toctree: api
+
+    ~disable
+    ~enable
+    ~enabled
+    ~create_operator_primitive
+    ~create_measurement_obs_primitive
+    ~create_measurement_wires_primitive
+    ~create_measurement_mcm_primitive
+    ~qnode_call
 
 To activate and deactivate the new PennyLane program capturing mechanism, use
 the switches ``qml.capture.enable`` and ``qml.capture.disable``.
@@ -114,4 +127,18 @@ def _(*args, **kwargs):
 """
 from .switches import disable, enable, enabled
 from .capture_meta import CaptureMeta
-from .primitives import create_operator_primitive
+from .primitives import (
+    create_operator_primitive,
+    create_measurement_obs_primitive,
+    create_measurement_wires_primitive,
+    create_measurement_mcm_primitive,
+)
+from .capture_qnode import qnode_call
+
+
+def __getattr__(key):
+    if key == "AbstractOperator":
+        from .primitives import _get_abstract_operator  # pylint: disable=import-outside-toplevel
+
+        return _get_abstract_operator()
+    raise AttributeError(f"module 'pennylane.capture' has no attribute '{key}'")

pennylane/capture/capture_meta.py

@@ -28,6 +28,39 @@ class CaptureMeta(type):
 
     See ``pennylane/capture/explanations.md`` for more detailed information on how this technically
     works.
+
+    .. code-block::
+
+        class AbstractMyObj(jax.core.AbstractValue):
+            pass
+
+        jax.core.raise_to_shaped_mappings[AbstractMyObj] = lambda aval, _: aval
+
+        class MyObj(metaclass=qml.capture.CaptureMeta):
+
+            primitive = jax.core.Primitive("MyObj")
+
+            @classmethod
+            def _primitive_bind_call(cls, a):
+                return cls.primitive.bind(a)
+
+            def __init__(self, a):
+                self.a = a
+
+        @MyObj.primitive.def_impl
+        def _(a):
+            return type.__call__(MyObj, a)
+
+        @MyObj.primitive.def_abstract_eval
+        def _(a):
+            return AbstractMyObj()
+
+    >>> jaxpr = jax.make_jaxpr(MyObj)(0.1)
+    >>> jaxpr
+    { lambda ; a:f32[]. let b:AbstractMyObj() = MyObj a in (b,) }
+    >>> jax.core.eval_jaxpr(jaxpr.jaxpr, jaxpr.consts, 0.1)
+    [<__main__.MyObj at 0x17fc3ea50>]
+
     """
 
     @property

pennylane/capture/capture_qnode.py

@@ -0,0 +1,162 @@
+# 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.
+"""
+This submodule defines a capture compatible call to QNodes.
+"""
+
+from functools import lru_cache, partial
+
+import pennylane as qml
+
+has_jax = True
+try:
+    import jax
+except ImportError:
+    has_jax = False
+
+
+def _get_shapes_for(*measurements, shots=None, num_device_wires=0):
+    if jax.config.jax_enable_x64:
+        dtype_map = {
+            float: jax.numpy.float64,
+            int: jax.numpy.int64,
+            complex: jax.numpy.complex128,
+        }
+    else:
+        dtype_map = {
+            float: jax.numpy.float32,
+            int: jax.numpy.int32,
+            complex: jax.numpy.complex64,
+        }
+
+    shapes = []
+    if not shots:
+        shots = [None]
+
+    for s in shots:
+        for m in measurements:
+            shape, dtype = m.abstract_eval(shots=s, num_device_wires=num_device_wires)
+            shapes.append(jax.core.ShapedArray(shape, dtype_map.get(dtype, dtype)))
+    return shapes
+
+
+@lru_cache()
+def _get_qnode_prim():
+    if not has_jax:
+        return None
+    qnode_prim = jax.core.Primitive("qnode")
+    qnode_prim.multiple_results = True
+
+    @qnode_prim.def_impl
+    def _(*args, shots, device, qnode_kwargs, qfunc_jaxpr):
+        def qfunc(*inner_args):
+            return jax.core.eval_jaxpr(qfunc_jaxpr.jaxpr, qfunc_jaxpr.consts, *inner_args)
+
+        qnode = qml.QNode(qfunc, device, **qnode_kwargs)
+        return qnode._impl_call(*args, shots=shots)  # pylint: disable=protected-access
+
+    # pylint: disable=unused-argument
+    @qnode_prim.def_abstract_eval
+    def _(*args, shots, device, qnode_kwargs, qfunc_jaxpr):
+        mps = qfunc_jaxpr.out_avals
+        return _get_shapes_for(*mps, shots=shots, num_device_wires=len(device.wires))
+
+    return qnode_prim
+
+
+# pylint: disable=protected-access
+def _get_device_shots(device) -> "qml.measurements.Shots":
+    if isinstance(device, qml.devices.LegacyDevice):
+        if device._shot_vector:
+            return qml.measurements.Shots(device._raw_shot_sequence)
+        return qml.measurements.Shots(device.shots)
+    return device.shots
+
+
+def qnode_call(qnode: "qml.QNode", *args, **kwargs) -> "qml.typing.Result":
+    """A capture compatible call to a QNode. This function is internally used by ``QNode.__call__``.
+
+    Args:
+        qnode (QNode): a QNode
+        args: the arguments the QNode is called with
+
+    Keyword Args:
+        Any keyword arguments accepted by the quantum function
+
+    Returns:
+        qml.typing.Result: the result of a qnode execution
+
+    **Example:**
+
+    .. code-block:: python
+
+        @qml.qnode(qml.device('lightning.qubit', wires=1))
+        def circuit(x):
+            qml.RX(x, wires=0)
+            return qml.expval(qml.Z(0)), qml.probs()
+
+        def f(x):
+            expval_z, probs = circuit(np.pi * x, shots=50)
+            return 2*expval_z + probs
+
+        jaxpr = jax.make_jaxpr(f)(0.1)
+        print("jaxpr: \n", jaxpr)
+
+        res = jax.core.eval_jaxpr(jaxpr.jaxpr, jaxpr.consts, 0.7)
+        print("\nresult: \n", res)
+
+
+    .. code-block:: none
+
+        jaxpr:
+        { lambda ; a:f32[]. let
+            b:f32[] = mul 3.141592653589793 a
+            c:f32[] d:f32[2] = qnode[
+            device=<lightning.qubit device (wires=1) at 0x30755b3d0>
+            qfunc_jaxpr={ lambda ; e:f32[]. let
+                _:AbstractOperator() = RX[n_wires=1] e 0
+                f:AbstractOperator() = PauliZ[n_wires=1] 0
+                g:AbstractMeasurement(n_wires=None) = expval_obs f
+                h:AbstractMeasurement(n_wires=0) = probs_wires
+                in (g, h) }
+            qnode_kwargs={'diff_method': 'best', 'grad_on_execution': 'best', 'cache': False, 'cachesize': 10000, 'max_diff': 1, 'max_expansion': 10, 'device_vjp': False}
+            shots=Shots(total=50)
+            ] b
+            i:f32[] = mul 2.0 c
+            j:f32[2] = add i d
+        in (j,) }
+
+        result:
+        [Array([-1.08, -0.64], dtype=float32)]
+
+
+    """
+    shots = kwargs.pop("shots", _get_device_shots(qnode.device))
+    shots = qml.measurements.Shots(shots)
+    if shots.has_partitioned_shots:
+        # Questions over the pytrees and the nested result object shape
+        raise NotImplementedError("shot vectors are not yet supported with plxpr capture.")
+
+    if not qnode.device.wires:
+        raise NotImplementedError("devices must specify wires for integration with plxpr capture.")
+
+    qfunc = partial(qnode.func, **kwargs) if kwargs else qnode.func
+
+    qfunc_jaxpr = jax.make_jaxpr(qfunc)(*args)
+    qnode_kwargs = {"diff_method": qnode.diff_method, **qnode.execute_kwargs}
+    qnode_prim = _get_qnode_prim()
+
+    return qnode_prim.bind(
+        *args, shots=shots, device=qnode.device, qnode_kwargs=qnode_kwargs, qfunc_jaxpr=qfunc_jaxpr
+    )