[Program Capture] Capture & execute qml.jacobian in plxpr (#6127)

**Context:**
We're adding support for differentiation in plxpr, also see #6120.

**Description of the Change:**
This PR adds support for `qml.jacobian`, similar to the support for
`qml.grad`.
Note that Pytree support will be needed to allow for multi-argument
derivatives.

**Benefits:**
Capture derivatives of non-scalar functions.

**Possible Drawbacks:**
See discussion around `qml.grad` in #6120.

**Related GitHub Issues:**

[sc-71860]

---------

Co-authored-by: Christina Lee <[email protected]>

doc/releases/changelog-dev.md

@@ -20,6 +20,15 @@
   which differs from the Autograd implementation of `qml.grad` itself.
   [(#6120)](https://github.com/PennyLaneAI/pennylane/pull/6120)
 
+<h4>Capturing and representing hybrid programs</h4>
+
+* Differentiation of hybrid programs via `qml.grad` and `qml.jacobian` can now be captured
+  into plxpr. When evaluating a captured `qml.grad` (`qml.jacobian`) instruction, it will
+  dispatch to `jax.grad` (`jax.jacobian`), which differs from the Autograd implementation
+  without capture.
+  [(#6120)](https://github.com/PennyLaneAI/pennylane/pull/6120)
+  [(#6127)](https://github.com/PennyLaneAI/pennylane/pull/6127)
+
 * Improve unit testing for capturing of nested control flows.
   [(#6111)](https://github.com/PennyLaneAI/pennylane/pull/6111)
 

pennylane/_grad.py

@@ -24,7 +24,7 @@
 from autograd.wrap_util import unary_to_nary
 
 from pennylane.capture import enabled
-from pennylane.capture.capture_diff import _get_grad_prim
+from pennylane.capture.capture_diff import _get_grad_prim, _get_jacobian_prim
 from pennylane.compiler import compiler
 from pennylane.compiler.compiler import CompileError
 
@@ -434,8 +434,11 @@ def circuit(x):
         ops_loader = available_eps[active_jit]["ops"].load()
         return ops_loader.jacobian(func, method=method, h=h, argnums=argnum)
 
+    if enabled():
+        return _capture_diff(func, argnum, _get_jacobian_prim(), method=method, h=h)
+
     if method or h:
-        raise ValueError(f"Invalid values for 'method={method}' and 'h={h}' in interpreted mode")
+        raise ValueError(f"Invalid values '{method=}' and '{h=}' without QJIT.")
 
     def _get_argnum(args):
         """Inspect the arguments for differentiability and return the

pennylane/capture/capture_diff.py

@@ -82,3 +82,38 @@ def _(*args, argnum, jaxpr, n_consts, method, h):
         return tuple(jaxpr.invars[i].aval for i in argnum)
 
     return grad_prim
+
+
+@lru_cache
+def _get_jacobian_prim():
+    """Create a primitive for Jacobian computations.
+    This primitive is used when capturing ``qml.jacobian``.
+    """
+    jacobian_prim = create_non_jvp_primitive()("jacobian")
+    jacobian_prim.multiple_results = True  # pylint: disable=attribute-defined-outside-init
+
+    # pylint: disable=too-many-arguments
+    @jacobian_prim.def_impl
+    def _(*args, argnum, jaxpr, n_consts, method, h):
+        if method or h:  # pragma: no cover
+            raise ValueError(f"Invalid values '{method=}' and '{h=}' without QJIT.")
+        consts = args[:n_consts]
+        args = args[n_consts:]
+
+        def func(*inner_args):
+            return jax.core.eval_jaxpr(jaxpr, consts, *inner_args)
+
+        return jax.jacobian(func, argnums=argnum)(*args)
+
+    # pylint: disable=unused-argument
+    @jacobian_prim.def_abstract_eval
+    def _(*args, argnum, jaxpr, n_consts, method, h):
+        in_avals = [jaxpr.invars[i].aval for i in argnum]
+        out_shapes = (outvar.aval.shape for outvar in jaxpr.outvars)
+        return [
+            jax.core.ShapedArray(out_shape + in_aval.shape, in_aval.dtype)
+            for out_shape in out_shapes
+            for in_aval in in_avals
+        ]
+
+    return jacobian_prim

pennylane/capture/primitives.py

@@ -22,7 +22,7 @@
 from pennylane.ops.op_math.condition import _get_cond_qfunc_prim
 from pennylane.ops.op_math.controlled import _get_ctrl_qfunc_prim
 
-from .capture_diff import _get_grad_prim
+from .capture_diff import _get_grad_prim, _get_jacobian_prim
 from .capture_measurements import _get_abstract_measurement
 from .capture_operators import _get_abstract_operator
 from .capture_qnode import _get_qnode_prim
@@ -32,6 +32,7 @@
 adjoint_transform_prim = _get_adjoint_qfunc_prim()
 ctrl_transform_prim = _get_ctrl_qfunc_prim()
 grad_prim = _get_grad_prim()
+jacobian_prim = _get_jacobian_prim()
 qnode_prim = _get_qnode_prim()
 cond_prim = _get_cond_qfunc_prim()
 for_loop_prim = _get_for_loop_qfunc_prim()
@@ -44,6 +45,7 @@
     "adjoint_transform_prim",
     "ctrl_transform_prim",
     "grad_prim",
+    "jacobian_prim",
     "qnode_prim",
     "cond_prim",
     "for_loop_prim",