[Capture] handle grad and jacobian (#2078)

**Context:**

Handles the `grad` primitive from pennylane and adds a lowering for it.

This requires PennyLane PR:
PennyLaneAI/pennylane#8357

Once that merges in, we can retarget this branch against the newer
pennylane.

**Description of the Change:**

Instead of translating the pennylane grad prim to the catalyst one, we
just register a lowering for the grad primitive. We also have to
register a translation rule where we are translating the target jaxpr.

**Benefits:**

Can handle grad and jacobian calls.

**Possible Drawbacks:**

**Related GitHub Issues:**

[sc-100560]

Author: albi3ro

.dep-versions

@@ -10,7 +10,7 @@ enzyme=v0.0.186
 
 # For a custom PL version, update the package version here and at
 # 'doc/requirements.txt'
-pennylane=0.44.0.dev17
+pennylane=0.44.0.dev20
 
 # For a custom LQ/LK version, update the package version here and at
 # 'doc/requirements.txt'

doc/releases/changelog-dev.md

@@ -20,6 +20,9 @@
 * Pytree inputs can now be used when program capture is enabled.
   [(#2165)](https://github.com/PennyLaneAI/catalyst/pull/2165)
 
+* `qml.grad` and `qml.jacobian` can now be used with `qjit` when program capture is enabled.
+  [(#2078)](https://github.com/PennyLaneAI/catalyst/pull/2078)
+
 <h3>Breaking changes 💔</h3>
 
 <h3>Deprecations 👋</h3>

doc/requirements.txt

@@ -33,4 +33,4 @@ lxml_html_clean
 --extra-index-url https://test.pypi.org/simple/
 pennylane-lightning-kokkos==0.43.0
 pennylane-lightning==0.43.0
-pennylane==0.44.0.dev17
+pennylane==0.44.0.dev20

frontend/catalyst/from_plxpr/from_plxpr.py

@@ -16,7 +16,6 @@
 """
 # pylint: disable=protected-access
 
-
 import warnings
 from copy import copy
 from functools import partial
@@ -28,6 +27,7 @@
 from jax.extend.linear_util import wrap_init
 from pennylane.capture import PlxprInterpreter, qnode_prim
 from pennylane.capture.expand_transforms import ExpandTransformsInterpreter
+from pennylane.capture.primitives import jacobian_prim as pl_jac_prim
 from pennylane.ops.functions.map_wires import _map_wires_transform as pl_map_wires
 from pennylane.transforms import cancel_inverses as pl_cancel_inverses
 from pennylane.transforms import commute_controlled as pl_commute_controlled
@@ -157,6 +157,18 @@ def __init__(self):
         super().__init__()
 
 
+@WorkflowInterpreter.register_primitive(pl_jac_prim)
+def handle_grad(self, *args, jaxpr, n_consts, **kwargs):
+    """Translate a grad equation."""
+    f = partial(copy(self).eval, jaxpr, args[:n_consts])
+    new_jaxpr = jax.make_jaxpr(f)(*args[n_consts:])
+
+    new_args = (*new_jaxpr.consts, *args[n_consts:])
+    return pl_jac_prim.bind(
+        *new_args, jaxpr=new_jaxpr.jaxpr, n_consts=len(new_jaxpr.consts), **kwargs
+    )
+
+
 # pylint: disable=unused-argument, too-many-arguments
 @WorkflowInterpreter.register_primitive(qnode_prim)
 def handle_qnode(

frontend/catalyst/jax_primitives.py

@@ -100,6 +100,7 @@
     VarianceOp,
 )
 from mlir_quantum.dialects.quantum import YieldOp as QYieldOp
+from pennylane.capture.primitives import jacobian_prim as pl_jac_prim
 
 from catalyst.compiler import get_lib_path
 from catalyst.jax_extras import (
@@ -644,9 +645,12 @@ def _grad_lowering(ctx, *args, jaxpr, fn, grad_params):
     consts = []
     offset = len(args) - len(jaxpr.consts)
     for i, jax_array_or_tracer in enumerate(jaxpr.consts):
-        if not isinstance(
-            jax_array_or_tracer, jax._src.interpreters.partial_eval.DynamicJaxprTracer
-        ):
+        if isinstance(jax_array_or_tracer, jax._src.interpreters.partial_eval.DynamicJaxprTracer):
+            # There are some cases where this value cannot be converted into
+            # a jax.numpy.array.
+            # in that case we get it from the arguments.
+            consts.append(args[offset + i])
+        else:
             # ``ir.DenseElementsAttr.get()`` constructs a dense elements attribute from an array of
             # element values. This doesn't support ``jaxlib.xla_extension.Array``, so we have to
             # cast such constants to numpy array types.
@@ -656,11 +660,6 @@ def _grad_lowering(ctx, *args, jaxpr, fn, grad_params):
             attr = ir.DenseElementsAttr.get(nparray, type=const_type)
             constval = StableHLOConstantOp(attr).results
             consts.append(constval)
-        else:
-            # There are some cases where this value cannot be converted into
-            # a jax.numpy.array.
-            # in that case we get it from the arguments.
-            consts.append(args[offset + i])
 
     method, h, argnums = grad_params.method, grad_params.h, grad_params.expanded_argnums
     mlir_ctx = ctx.module_context.context
@@ -673,7 +672,6 @@ def _grad_lowering(ctx, *args, jaxpr, fn, grad_params):
     argnum_numpy = np.array(new_argnums)
     diffArgIndices = ir.DenseIntElementsAttr.get(argnum_numpy)
     func_op = lower_jaxpr(ctx, jaxpr, (method, h, *argnums))
-
     symbol_ref = get_symbolref(ctx, func_op)
     output_types = list(map(mlir.aval_to_ir_types, ctx.avals_out))
     flat_output_types = util.flatten(output_types)
@@ -692,6 +690,30 @@ def _grad_lowering(ctx, *args, jaxpr, fn, grad_params):
     ).results
 
 
+# pylint: disable=too-many-arguments
+def _capture_grad_lowering(ctx, *args, argnums, jaxpr, n_consts, method, h, fn, scalar_out):
+    mlir_ctx = ctx.module_context.context
+    f64 = ir.F64Type.get(mlir_ctx)
+    finiteDiffParam = ir.FloatAttr.get(f64, h)
+
+    new_argnums = [num + n_consts for num in argnums]
+    argnum_numpy = np.array(new_argnums)
+    diffArgIndices = ir.DenseIntElementsAttr.get(argnum_numpy)
+    func_op = lower_jaxpr(ctx, jaxpr, (method, h, *new_argnums), fn=fn)
+    symbol_ref = get_symbolref(ctx, func_op)
+    output_types = list(map(mlir.aval_to_ir_types, ctx.avals_out))
+    flat_output_types = util.flatten(output_types)
+
+    return GradOp(
+        flat_output_types,
+        ir.StringAttr.get(method),
+        symbol_ref,
+        mlir.flatten_lowering_ir_args(args),
+        diffArgIndices=diffArgIndices,
+        finiteDiffParam=finiteDiffParam,
+    ).results
+
+
 # value_and_grad
 #
 @value_and_grad_p.def_impl
@@ -2542,6 +2564,7 @@ def subroutine_lowering(*args, **kwargs):
     (while_p, _while_loop_lowering),
     (for_p, _for_loop_lowering),
     (grad_p, _grad_lowering),
+    (pl_jac_prim, _capture_grad_lowering),
     (func_p, _func_lowering),
     (jvp_p, _jvp_lowering),
     (vjp_p, _vjp_lowering),

frontend/catalyst/jax_primitives_utils.py

@@ -29,6 +29,31 @@
 from catalyst.jax_extras.lowering import get_mlir_attribute_from_pyval
 
 
+def _only_single_expval(call_jaxpr: core.ClosedJaxpr) -> bool:
+    found_expval = False
+    for eqn in call_jaxpr.eqns:
+        name = eqn.primitive.name
+        if name in {"probs", "counts", "sample"}:
+            return False
+        elif name == "expval":
+            if found_expval:
+                return False
+            found_expval = True
+    return True
+
+
+def _calculate_diff_method(qn: qml.QNode, call_jaxpr: core.ClosedJaxpr):
+    diff_method = str(qn.diff_method)
+    if diff_method != "best":
+        return diff_method
+
+    device_name = getattr(getattr(qn, "device", None), "name", None)
+
+    if device_name and "lightning" in device_name and _only_single_expval(call_jaxpr):
+        return "adjoint"
+    return "parameter-shift"
+
+
 def get_call_jaxpr(jaxpr):
     """Extracts the `call_jaxpr` from a JAXPR if it exists.""" ""
     for eqn in jaxpr.eqns:
@@ -45,28 +70,36 @@ def get_call_equation(jaxpr):
     raise AssertionError("No call_jaxpr found in the JAXPR.")
 
 
-def lower_jaxpr(ctx, jaxpr, context=None):
+def lower_jaxpr(ctx, jaxpr, metadata=None, fn=None):
     """Lowers a call primitive jaxpr, may be either func_p or quantum_kernel_p
 
     Args:
         ctx: LoweringRuleContext
         jaxpr: JAXPR to be lowered
-        context: additional context to distinguish different FuncOps
+        metadata: additional metadata to distinguish different FuncOps
+        fn (Callable | None): the function the jaxpr corresponds to. Used for naming and caching.
 
     Returns:
         FuncOp
     """
-    equation = get_call_equation(jaxpr)
-    call_jaxpr = equation.params["call_jaxpr"]
-    callable_ = equation.params.get("fn")
-    if callable_ is None:
-        callable_ = equation.params.get("qnode")
-    pipeline = equation.params.get("pipeline")
-    return lower_callable(ctx, callable_, call_jaxpr, pipeline=pipeline, context=context)
+
+    if fn is None or isinstance(fn, qml.QNode):
+        equation = get_call_equation(jaxpr)
+        call_jaxpr = equation.params["call_jaxpr"]
+        pipeline = equation.params.get("pipeline")
+        callable_ = equation.params.get("fn")
+        if callable_ is None:
+            callable_ = equation.params.get("qnode", None)
+    else:
+        call_jaxpr = jaxpr
+        pipeline = ()
+        callable_ = fn
+
+    return lower_callable(ctx, callable_, call_jaxpr, pipeline=pipeline, metadata=metadata)
 
 
 # pylint: disable=too-many-arguments, too-many-positional-arguments
-def lower_callable(ctx, callable_, call_jaxpr, pipeline=None, context=None, public=False):
+def lower_callable(ctx, callable_, call_jaxpr, pipeline=(), metadata=None, public=False):
     """Lowers _callable to MLIR.
 
     If callable_ is a qnode, then we will first create a module, then
@@ -86,33 +119,33 @@ def lower_callable(ctx, callable_, call_jaxpr, pipeline=None, context=None, publ
     if pipeline is None:
         pipeline = tuple()
 
-    if not isinstance(callable_, qml.QNode):
-        return get_or_create_funcop(
-            ctx, callable_, call_jaxpr, pipeline, context=context, public=public
-        )
-
-    return get_or_create_qnode_funcop(ctx, callable_, call_jaxpr, pipeline, context=context)
+    if isinstance(callable_, qml.QNode):
+        return get_or_create_qnode_funcop(ctx, callable_, call_jaxpr, pipeline, metadata=metadata)
+    return get_or_create_funcop(
+        ctx, callable_, call_jaxpr, pipeline, metadata=metadata, public=public
+    )
 
 
 # pylint: disable=too-many-arguments, too-many-positional-arguments
-def get_or_create_funcop(ctx, callable_, call_jaxpr, pipeline, context=None, public=False):
+def get_or_create_funcop(ctx, callable_, call_jaxpr, pipeline, metadata=None, public=False):
     """Get funcOp from cache, or create it from scratch
 
     Args:
         ctx: LoweringRuleContext
         callable_: python function
         call_jaxpr: jaxpr representing callable_
-        context: additional context to distinguish different FuncOps
+        metadata: additional metadata to distinguish different FuncOps
         public: whether the visibility should be marked public
 
     Returns:
         FuncOp
     """
-    if context is None:
-        context = tuple()
-    key = (callable_, *context, *pipeline)
-    if func_op := get_cached(ctx, key):
-        return func_op
+    if metadata is None:
+        metadata = tuple()
+    key = (callable_, *metadata, *pipeline)
+    if callable_ is not None:
+        if func_op := get_cached(ctx, key):
+            return func_op
     func_op = lower_callable_to_funcop(ctx, callable_, call_jaxpr, public=public)
     cache(ctx, key, func_op)
     return func_op
@@ -135,10 +168,10 @@ def lower_callable_to_funcop(ctx, callable_, call_jaxpr, public=False):
 
     kwargs = {}
     kwargs["ctx"] = ctx.module_context
-    if not isinstance(callable_, functools.partial):
-        name = callable_.__name__
-    else:
+    if isinstance(callable_, functools.partial):
         name = callable_.func.__name__ + ".partial"
+    else:
+        name = callable_.__name__
 
     kwargs["name"] = name
     kwargs["jaxpr"] = call_jaxpr
@@ -154,28 +187,7 @@ def lower_callable_to_funcop(ctx, callable_, call_jaxpr, public=False):
     if isinstance(callable_, qml.QNode):
         func_op.attributes["qnode"] = ir.UnitAttr.get()
 
-        diff_method = str(callable_.diff_method)
-
-        if diff_method == "best":
-
-            def only_single_expval():
-                found_expval = False
-                for eqn in call_jaxpr.eqns:
-                    name = eqn.primitive.name
-                    if name in {"probs", "counts", "sample"}:
-                        return False
-                    elif name == "expval":
-                        if found_expval:
-                            return False
-                        found_expval = True
-                return True
-
-            device_name = getattr(getattr(callable_, "device", None), "name", None)
-
-            if device_name and "lightning" in device_name and only_single_expval():
-                diff_method = "adjoint"
-            else:
-                diff_method = "parameter-shift"
+        diff_method = _calculate_diff_method(callable_, call_jaxpr)
 
         func_op.attributes["diff_method"] = ir.StringAttr.get(diff_method)
 
@@ -195,7 +207,7 @@ def only_single_expval():
     return func_op
 
 
-def get_or_create_qnode_funcop(ctx, callable_, call_jaxpr, pipeline, context):
+def get_or_create_qnode_funcop(ctx, callable_, call_jaxpr, pipeline, metadata):
     """A wrapper around lower_qnode_to_funcop that will cache the FuncOp.
 
     Args:
@@ -205,11 +217,11 @@ def get_or_create_qnode_funcop(ctx, callable_, call_jaxpr, pipeline, context):
     Returns:
       FuncOp
     """
-    if context is None:
-        context = tuple()
+    if metadata is None:
+        metadata = tuple()
     if callable_.static_argnums:
         return lower_qnode_to_funcop(ctx, callable_, call_jaxpr, pipeline)
-    key = (callable_, *context, *pipeline)
+    key = (callable_, *metadata, *pipeline)
     if func_op := get_cached(ctx, key):
         return func_op
     func_op = lower_qnode_to_funcop(ctx, callable_, call_jaxpr, pipeline)