[Program Capture] Capture & execute qml.grad in plxpr

doc/releases/changelog-dev.md

@@ -150,6 +150,11 @@
 * `qml.for_loop` now supports `range`-like syntax with default `step=1`.
   [(#6068)](https://github.com/PennyLaneAI/pennylane/pull/6068)
 
+* Differentiation of hybrid programs via `qml.grad` can now be captured into plxpr.
+  When evaluating a captured `qml.grad` instruction, it will dispatch to `jax.grad`,
+  which differs from the Autograd implementation of `qml.grad` itself.
+  [(#6120)](https://github.com/PennyLaneAI/pennylane/pull/6120)
+
 * Applying `adjoint` and `ctrl` to a quantum function can now be captured into plxpr.
   Furthermore, the `qml.cond` function can be captured into plxpr.
   [(#5966)](https://github.com/PennyLaneAI/pennylane/pull/5966)

pennylane/_grad.py

@@ -15,19 +15,39 @@
 This module contains the autograd wrappers :class:`grad` and :func:`jacobian`
 """
 import warnings
+from functools import partial, wraps
 
 from autograd import jacobian as _jacobian
 from autograd.core import make_vjp as _make_vjp
 from autograd.extend import vspace
 from autograd.numpy.numpy_boxes import ArrayBox
 from autograd.wrap_util import unary_to_nary
 
+from pennylane.capture import create_grad_primitive, enabled
 from pennylane.compiler import compiler
 from pennylane.compiler.compiler import CompileError
 
 make_vjp = unary_to_nary(_make_vjp)
 
 
+def _capture_diff(func, argnum=None, diff_prim=None):
+    """Capture-compatible gradient computation."""
+    import jax  # pylint: disable=import-outside-toplevel
+
+    if isinstance(argnum, int):
+        argnum = [argnum]
+    if argnum is None:
+        argnum = [0]
+
+    @wraps(func)
+    def new_func(*args, **kwargs):
+        jaxpr = jax.make_jaxpr(partial(func, **kwargs))(*args)
+        prim_kwargs = {"argnum": argnum, "jaxpr": jaxpr.jaxpr, "n_consts": len(jaxpr.consts)}
+        return diff_prim.bind(*jaxpr.consts, *args, **prim_kwargs)
+
+    return new_func
+
+
 class grad:
     """Returns the gradient as a callable function of hybrid quantum-classical functions.
     :func:`~.qjit` and Autograd compatible.
@@ -97,9 +117,10 @@ def __new__(cls, func, argnum=None, method=None, h=None):
             return ops_loader.grad(func, method=method, h=h, argnums=argnum)
 
         if method or h:  # pragma: no cover
-            raise ValueError(
-                f"Invalid values for 'method={method}' and 'h={h}' in interpreted mode"
-            )
+            raise ValueError(f"Invalid values '{method=}' and '{h=}' without QJIT.")
+
+        if enabled():
+            return _capture_diff(func, argnum, create_grad_primitive())
 
         return super().__new__(cls)
 

pennylane/capture/__init__.py

@@ -128,6 +128,7 @@ def _(*args, **kwargs):
 from .switches import disable, enable, enabled
 from .capture_meta import CaptureMeta, ABCCaptureMeta
 from .primitives import (
+    create_grad_primitive,
     create_operator_primitive,
     create_measurement_obs_primitive,
     create_measurement_wires_primitive,
@@ -172,6 +173,7 @@ def __getattr__(key):
     "create_measurement_obs_primitive",
     "create_measurement_wires_primitive",
     "create_measurement_mcm_primitive",
+    "create_grad_primitive",
     "qnode_call",
     "AbstractOperator",
     "AbstractMeasurement",

pennylane/capture/capture_qnode.py

@@ -27,6 +27,8 @@
 has_jax = True
 try:
     import jax
+    from jax.interpreters import ad
+
 except ImportError:
     has_jax = False
 
@@ -87,6 +89,11 @@ def _(*args, qnode, shots, device, qnode_kwargs, qfunc_jaxpr, n_consts):
         mps = qfunc_jaxpr.outvars
         return _get_shapes_for(*mps, shots=shots, num_device_wires=len(device.wires))
 
+    def _qnode_jvp(*args_and_tangents, **impl_kwargs):
+        return jax.jvp(partial(qnode_prim.impl, **impl_kwargs), *args_and_tangents)
+
+    ad.primitive_jvps[qnode_prim] = _qnode_jvp
+
     return qnode_prim
 
 

pennylane/capture/primitives.py

@@ -163,6 +163,31 @@ def __hash__(self):
     return AbstractMeasurement
 
 
+@lru_cache
+def create_non_jvp_primitive():
+    """Create a primitive type ``NonJVPPrimitive``, which binds to JAX's JVPTrace
+    like a standard Python function and otherwise behaves like jax.core.Primitive.
+    """
+
+    if not has_jax:  # pragma: no cover
+        return None
+
+    # pylint: disable=too-few-public-methods
+    class NonJVPPrimitive(jax.core.Primitive):
+        """A subclass to JAX's Primitive that works like a Python function
+        when evaluating JVPTracers."""
+
+        def bind_with_trace(self, trace, args, params):
+            """Bind the ``NonJVPPrimitive`` with a trace. If the trace is a ``JVPTrace``,
+            binding falls back to a standard Python function call. Otherwise, the
+            bind call of JAX's standard Primitive is used."""
+            if isinstance(trace, jax.interpreters.ad.JVPTrace):
+                return self.impl(*args, **params)
+            return super().bind_with_trace(trace, args, params)
+
+    return NonJVPPrimitive
+
+
 def create_operator_primitive(
     operator_type: Type["qml.operation.Operator"],
 ) -> Optional["jax.core.Primitive"]:
@@ -182,7 +207,7 @@ def create_operator_primitive(
     if not has_jax:
         return None
 
-    primitive = jax.core.Primitive(operator_type.__name__)
+    primitive = create_non_jvp_primitive()(operator_type.__name__)
 
     @primitive.def_impl
     def _(*args, **kwargs):
@@ -318,3 +343,32 @@ def _(*args, has_eigvals=False, **_):
         return abstract_type(abstract_eval, n_wires=n_wires, has_eigvals=has_eigvals)
 
     return primitive
+
+
+@lru_cache
+def create_grad_primitive():
+    """Create a primitive for gradient computations.
+    This primitive is used when capturing ``qml.grad``.
+    """
+    grad_prim = create_non_jvp_primitive()("grad")
+    grad_prim.multiple_results = True  # pylint: disable=attribute-defined-outside-init
+
+    # pylint: disable=too-many-arguments
+    @grad_prim.def_impl
+    def _(*args, argnum, jaxpr, n_consts):
+        consts = args[:n_consts]
+        args = args[n_consts:]
+
+        def func(*inner_args):
+            return jax.core.eval_jaxpr(jaxpr, consts, *inner_args)[0]
+
+        return jax.grad(func, argnums=argnum)(*args)
+
+    # pylint: disable=unused-argument
+    @grad_prim.def_abstract_eval
+    def _(*args, argnum, jaxpr, n_consts):
+        if len(jaxpr.outvars) != 1 or jaxpr.outvars[0].aval.shape != ():
+            raise TypeError("Grad only applies to scalar-output functions. Try jacobian or egrad.")
+        return tuple(jaxpr.invars[i].aval for i in argnum)
+
+    return grad_prim