Use grad context for hashing the generated stablehlo program (#1604)

**Context:** After PR #1562, a single function could have multiple JAXPR
representations based on whether it was under a grad context or not.
This made the previous hash based on the function id create possible
conflicts. To address this, we hashed on the jaxpr string
representation. (We cannot hash on the jax object itself since they are
unique).

The JAXPR string representation can be very long and hashing over long
strings can take a long time.

**Description of the Change:** Instead of hashing the string
representation, add a simple key to denote whether it is inside a grad
context or not.

**Benefits:** Reduced compilation time.

**Possible Drawbacks:** The cache key is getting more complicated. Maybe
the drawbacks outweight the benefits now?

**Related GitHub Issues:**

[sc-88454]

Author: erick-xanadu

doc/releases/changelog-0.11.0.md

@@ -193,6 +193,7 @@
   [(#1562)](https://github.com/PennyLaneAI/catalyst/pull/1562)
   [(#1568)](https://github.com/PennyLaneAI/catalyst/pull/1568)
   [(#1569)](https://github.com/PennyLaneAI/catalyst/pull/1569)
+  [(#1604)](https://github.com/PennyLaneAI/catalyst/pull/1604)
 
   Gates that are constant, such as when all parameters are Python or NumPy data types, are not
   decomposed when this is allowable. For the adjoint differentiation method, this is allowable

frontend/catalyst/jax_primitives.py

@@ -510,7 +510,7 @@ def _grad_lowering(ctx, *args, jaxpr, fn, grad_params):
     new_argnums = [num + offset for num in argnums]
     argnum_numpy = np.array(new_argnums)
     diffArgIndices = ir.DenseIntElementsAttr.get(argnum_numpy)
-    func_op = lower_jaxpr(ctx, jaxpr)
+    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))
@@ -585,7 +585,7 @@ def _value_and_grad_lowering(ctx, *args, jaxpr, fn, grad_params):
     val_result_types = flat_output_types[: len(flat_output_types) - len(argnums)]
     gradient_result_types = flat_output_types[len(flat_output_types) - len(argnums) :]
 
-    func_op = lower_jaxpr(ctx, jaxpr)
+    func_op = lower_jaxpr(ctx, jaxpr, (method, h, *argnums))
 
     symbol_ref = get_symbolref(ctx, func_op)
     return ValueAndGradOp(
@@ -635,7 +635,7 @@ def _jvp_lowering(ctx, *args, jaxpr, fn, grad_params):
     func_args = consts_and_args[: len(func_call_jaxpr.invars)]
     tang_args = consts_and_args[len(func_call_jaxpr.invars) :]
 
-    func_op = lower_jaxpr(ctx, jaxpr)
+    func_op = lower_jaxpr(ctx, jaxpr, (method, h, *argnums))
 
     assert (
         len(flat_output_types) % 2 == 0
@@ -688,7 +688,7 @@ def _vjp_lowering(ctx, *args, jaxpr, fn, grad_params):
     func_result_types = flat_output_types[: len(flat_output_types) - len(argnums)]
     vjp_result_types = flat_output_types[len(flat_output_types) - len(argnums) :]
 
-    func_op = lower_jaxpr(ctx, jaxpr)
+    func_op = lower_jaxpr(ctx, jaxpr, (method, h, *argnums))
 
     symbol_ref = get_symbolref(ctx, func_op)
     return VJPOp(

frontend/catalyst/jax_primitives_utils.py

@@ -47,18 +47,18 @@ def get_call_equation(jaxpr):
     raise AssertionError("No call_jaxpr found in the JAXPR.")
 
 
-def lower_jaxpr(ctx, jaxpr):
+def lower_jaxpr(ctx, jaxpr, context=None):
     """Lowers a call primitive jaxpr, may be either func_p or quantum_kernel_p"""
     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)
+    return lower_callable(ctx, callable_, call_jaxpr, pipeline=pipeline, context=context)
 
 
-def lower_callable(ctx, callable_, call_jaxpr, pipeline=None):
+def lower_callable(ctx, callable_, call_jaxpr, pipeline=None, context=None):
     """Lowers _callable to MLIR.
 
     If callable_ is a qnode, then we will first create a module, then
@@ -77,14 +77,16 @@ def lower_callable(ctx, callable_, call_jaxpr, pipeline=None):
         pipeline = tuple()
 
     if not isinstance(callable_, qml.QNode):
-        return get_or_create_funcop(ctx, callable_, call_jaxpr, pipeline)
+        return get_or_create_funcop(ctx, callable_, call_jaxpr, pipeline, context=context)
 
-    return get_or_create_qnode_funcop(ctx, callable_, call_jaxpr, pipeline)
+    return get_or_create_qnode_funcop(ctx, callable_, call_jaxpr, pipeline, context=context)
 
 
-def get_or_create_funcop(ctx, callable_, call_jaxpr, pipeline):
+def get_or_create_funcop(ctx, callable_, call_jaxpr, pipeline, context=None):
     """Get funcOp from cache, or create it from scratch"""
-    key = (str(call_jaxpr), *pipeline)
+    if context is None:
+        context = tuple()
+    key = (callable_, *context, *pipeline)
     if func_op := get_cached(ctx, key):
         return func_op
     func_op = lower_callable_to_funcop(ctx, callable_, call_jaxpr)
@@ -123,7 +125,7 @@ def lower_callable_to_funcop(ctx, callable_, call_jaxpr):
     return func_op
 
 
-def get_or_create_qnode_funcop(ctx, callable_, call_jaxpr, pipeline):
+def get_or_create_qnode_funcop(ctx, callable_, call_jaxpr, pipeline, context):
     """A wrapper around lower_qnode_to_funcop that will cache the FuncOp.
 
     Args:
@@ -133,7 +135,9 @@ def get_or_create_qnode_funcop(ctx, callable_, call_jaxpr, pipeline):
     Returns:
       FuncOp
     """
-    key = (str(call_jaxpr), *pipeline)
+    if context is None:
+        context = tuple()
+    key = (callable_, *context, *pipeline)
     if func_op := get_cached(ctx, key):
         return func_op
     func_op = lower_qnode_to_funcop(ctx, callable_, call_jaxpr, pipeline)