Capture a qnode into jaxpr

doc/releases/changelog-dev.md

@@ -108,9 +108,10 @@
   `m = measure(0); qml.sample(m)`.
   [(#5673)](https://github.com/PennyLaneAI/pennylane/pull/5673)
 
-* PennyLane operators and measurements can now automatically be captured as instructions in JAXPR.
+* 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)
   [(#5523)](https://github.com/PennyLaneAI/pennylane/pull/5523)
 
 * The `decompose` transform has an `error` kwarg to specify the type of error that should be raised, 

pennylane/capture/__init__.py

@@ -33,6 +33,7 @@
     ~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``.
@@ -132,6 +133,7 @@ def _(*args, **kwargs):
     create_measurement_wires_primitive,
     create_measurement_mcm_primitive,
 )
+from .capture_qnode import qnode_call
 
 
 def __getattr__(key):

pennylane/capture/capture_qnode.py

@@ -0,0 +1,167 @@
+# 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:")
+        print(jaxpr)
+
+        res = jax.core.eval_jaxpr(jaxpr.jaxpr, jaxpr.consts, 0.7)
+        print()
+        print("result:")
+        print(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 0x107639010>
+            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.3 , -0.74], dtype=float32)]
+
+
+    """
+    if "shots" in kwargs:
+        shots = qml.measurements.Shots(kwargs.pop("shots"))
+    else:
+        shots = _get_device_shots(qnode.device)
+    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
+    )

pennylane/workflow/qnode.py

@@ -1,4 +1,4 @@
-# Copyright 2018-2021 Xanadu Quantum Technologies Inc.
+# Copyright 2018-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.
@@ -1081,7 +1081,7 @@ def _execution_component(self, args: tuple, kwargs: dict, override_shots) -> qml
             res, self._qfunc_output, self._tape.shots.has_partitioned_shots
         )
 
-    def __call__(self, *args, **kwargs) -> qml.typing.Result:
+    def _impl_call(self, *args, **kwargs) -> qml.typing.Result:
 
         old_interface = self.interface
         if old_interface == "auto":
@@ -1113,6 +1113,11 @@ def __call__(self, *args, **kwargs) -> qml.typing.Result:
 
         return res
 
+    def __call__(self, *args, **kwargs) -> qml.typing.Result:
+        if qml.capture.enabled():
+            return qml.capture.qnode_call(self, *args, **kwargs)
+        return self._impl_call(*args, **kwargs)
+
 
 qnode = lambda device, **kwargs: functools.partial(QNode, device=device, **kwargs)
 qnode.__doc__ = QNode.__doc__