Update jax to 0.5.3 (#1652)

**Context:**
We update the jax version to 0.5.3.

For the detailed changes, see the corresponding PRs:

- LAPACK kernels are updated to adhere to the new JAX lowering rules for
external functions. #1685

- The trace stack is removed and replaced with a tracing context
manager. #1662

- A new `debug_info` argument/field is added to `Jaxpr`, the
`make_jaxpr` functions, and `jax.extend.linear_util.wrap_init`. #1670
#1671 #1681

[sc-88694]

---------

Co-authored-by: Erick Ochoa Lopez <[email protected]>
Co-authored-by: erick-xanadu <[email protected]>
Co-authored-by: Mehrdad Malek <[email protected]>
Co-authored-by: Mehrdad Malekmohammadi <[email protected]>

Author: 3 people

.dep-versions

@@ -1,5 +1,5 @@
 # Always update the version check in catalyst.__init__ when changing the JAX version.
-jax=0.4.28
+jax=0.5.3
 mhlo=89a891c986650c33df76885f5620e0a92150d90f
 llvm=3a8316216807d64a586b971f51695e23883331f7
 enzyme=v0.0.149

doc/releases/changelog-dev.md

@@ -76,6 +76,22 @@
   will handle direct decomposition of PPRs into PPMs.
   [(#1688)](https://github.com/PennyLaneAI/catalyst/pull/1688)
 
+* The version of JAX used by Catalyst is updated to 0.5.3.
+  [(#1652)](https://github.com/PennyLaneAI/catalyst/pull/1652)
+
+  Several internal changes were made for this update.
+    - LAPACK kernels are updated to adhere to the new JAX lowering rules for external functions.
+    [(#1685)](https://github.com/PennyLaneAI/catalyst/pull/1685)
+
+    - The trace stack is removed and replaced with a tracing context manager.
+    [(#1662)](https://github.com/PennyLaneAI/catalyst/pull/1662)
+
+    - A new `debug_info` argument is added to `Jaxpr`, the `make_jaxpr`
+    functions, and `jax.extend.linear_util.wrap_init`.
+    [(#1670)](https://github.com/PennyLaneAI/catalyst/pull/1670)
+    [(#1671)](https://github.com/PennyLaneAI/catalyst/pull/1671)
+    [(#1681)](https://github.com/PennyLaneAI/catalyst/pull/1681)
+
 <h3>Deprecations 👋</h3>
 
 <h3>Bug fixes 🐛</h3>
@@ -150,4 +166,5 @@ David Ittah,
 Tzung-Han Juang,
 Christina Lee,
 Erick Ochoa Lopez,
+Mehrdad Malekmohammadi,
 Paul Haochen Wang.

frontend/catalyst/__init__.py

@@ -23,7 +23,7 @@
 
 import jaxlib as _jaxlib
 
-_jaxlib_version = "0.4.28"
+_jaxlib_version = "0.5.3"
 if _jaxlib.__version__ != _jaxlib_version:
     import warnings
 

frontend/catalyst/api_extensions/callbacks.py

@@ -27,7 +27,7 @@
 
 import jax
 import jax.numpy as jnp
-from jax._src.api_util import shaped_abstractify
+from jax._src.core import shaped_abstractify
 from jax._src.tree_util import (
     Partial,
     tree_flatten,

frontend/catalyst/api_extensions/control_flow.py

@@ -24,9 +24,9 @@
 from typing import Callable, Iterable, List, Optional, Union
 
 import jax
-import jax.numpy as jnp
 from jax._src.api import _dtype
 from jax._src.tree_util import PyTreeDef, tree_flatten, tree_unflatten
+from jax.api_util import debug_info
 from pennylane import QNode
 
 import catalyst
@@ -46,6 +46,7 @@
 from catalyst.tracing.contexts import EvaluationContext, GradContext
 from catalyst.utils.callables import CatalystCallable
 from catalyst.utils.exceptions import DifferentiableCompileError
+from catalyst.utils.types import get_shape
 
 Differentiable = Union[Function, QNode]
 
@@ -474,10 +475,10 @@ def check_is_iterable(x, hint):
                     f"{_dtype(p)}, but got tangent dtype {_dtype(t)} instead."
                 )
 
-            if jnp.shape(p) != jnp.shape(t):
+            if get_shape(p) != get_shape(t):
                 raise ValueError(
                     "catalyst.jvp called with different function params and tangent shapes; "
-                    f"got function params shape {jnp.shape(p)} and tangent shape {jnp.shape(t)}"
+                    f"got function params shape {get_shape(p)} and tangent shape {get_shape(t)}"
                 )
 
         jaxpr, out_tree = _make_jaxpr_check_differentiable(fn, grad_params, *params)
@@ -584,11 +585,11 @@ def check_is_iterable(x, hint):
                     f"{_dtype(p)}, but got cotangent dtype {_dtype(t)} instead."
                 )
 
-            if jnp.shape(p) != jnp.shape(t):
+            if get_shape(p) != get_shape(t):
                 raise ValueError(
                     "catalyst.vjp called with different function output params and cotangent "
-                    f"shapes; got function output params shape {jnp.shape(p)} and cotangent shape "
-                    f"{jnp.shape(t)}"
+                    f"shapes; got function output params shape {get_shape(p)} and cotangent shape "
+                    f"{get_shape(t)}"
                 )
 
         cotangents, _ = tree_flatten(cotangents)
@@ -807,7 +808,9 @@ def _make_jaxpr_check_differentiable(
     return the output tree."""
     method = grad_params.method
     with mark_gradient_tracing(method):
-        jaxpr, _, out_tree = make_jaxpr2(f)(*args, **kwargs)
+        jaxpr, _, out_tree = make_jaxpr2(
+            f, debug_info=debug_info("grad make jaxpr", f, args, kwargs)
+        )(*args, **kwargs)
 
     for pos, arg in enumerate(jaxpr.in_avals):
         if arg.dtype.kind != "f" and pos in grad_params.expanded_argnums:

frontend/catalyst/api_extensions/differentiation.py

@@ -27,6 +27,7 @@
 import jax.numpy as jnp
 import numpy as np
 from jax._src.tree_util import tree_flatten, tree_leaves, tree_structure, tree_unflatten
+from jax.api_util import debug_info
 
 from catalyst.api_extensions.control_flow import for_loop
 from catalyst.jax_extras import make_jaxpr2
@@ -227,7 +228,9 @@ def __call__(self, *args, **kwargs):
         fn_args = tree_unflatten(args_tree, fn_args_flat)
 
         # Run 'fn' one time to get output-shape
-        _, shapes, init_result_tree = make_jaxpr2(self.fn)(*fn_args, **kwargs)
+        _, shapes, init_result_tree = make_jaxpr2(
+            self.fn, debug_info=debug_info("vmap", self.fn, args, kwargs)
+        )(*fn_args, **kwargs)
 
         init_result_flat = [jnp.zeros(shape=shape.shape, dtype=shape.dtype) for shape, _ in shapes]
         init_result = tree_unflatten(init_result_tree, init_result_flat)

frontend/catalyst/api_extensions/function_maps.py

@@ -26,6 +26,7 @@
 import jax
 import pennylane as qml
 from jax._src.tree_util import tree_flatten
+from jax.api_util import debug_info
 from jax.core import get_aval
 from pennylane import QueuingManager
 from pennylane.operation import Operator
@@ -146,7 +147,7 @@ def circuit():
     EvaluationContext.check_is_quantum_tracing(
         "catalyst.measure can only be used from within a qml.qnode."
     )
-    ctx = EvaluationContext.get_main_tracing_context()
+    cur_trace = EvaluationContext.get_current_trace()
     wires = list(wires) if isinstance(wires, (list, tuple)) else [wires]
     if len(wires) != 1:
         raise TypeError(f"Only one element is supported for the 'wires' parameter, got {wires}.")
@@ -161,7 +162,7 @@ def circuit():
     if postselect is not None and postselect not in [0, 1]:
         raise TypeError(f"postselect must be '0' or '1', got {postselect}")
 
-    m = new_inner_tracer(ctx.trace, get_aval(True))
+    m = new_inner_tracer(cur_trace, get_aval(True))
     MidCircuitMeasure(
         in_classical_tracers=in_classical_tracers,
         out_classical_tracers=[m],
@@ -344,14 +345,12 @@ def trace_quantum(self, ctx, device, trace, qrp, postselect_mode=None) -> QRegPr
         qubit = qrp.extract(self.wires)[0]
         if postselect_mode == "hw-like":
             qubit2 = self.bind_overwrite_classical_tracers(
-                ctx,
                 trace,
                 in_expanded_tracers=[qubit],
                 out_expanded_tracers=self.out_classical_tracers,
             )
         else:
             qubit2 = self.bind_overwrite_classical_tracers(
-                ctx,
                 trace,
                 in_expanded_tracers=[qubit],
                 out_expanded_tracers=self.out_classical_tracers,
@@ -405,12 +404,14 @@ def __call__(self, *args, **kwargs):
             return create_adjoint_op(base_op, self.lazy)
 
         tracing_artifacts = self.trace_body(args, kwargs)
-
-        return HybridAdjoint(*tracing_artifacts)
+        dbg = tracing_artifacts[3]
+        return HybridAdjoint(*tracing_artifacts[0:3], debug_info=dbg)
 
     def trace_body(self, args, kwargs):
         """Generate a HybridOpRegion for use by Catalyst."""
 
+        dbg = debug_info("AdjointCallable", self.target, args, kwargs)
+
         # Allow the creation of HybridAdjoint instances outside of any contexts.
         # Don't create a JAX context here as otherwise we could be dealing with escaped tracers.
         if not EvaluationContext.is_tracing():
@@ -420,18 +421,17 @@ def trace_body(self, args, kwargs):
 
             adjoint_region = HybridOpRegion(None, quantum_tape, [], [])
 
-            return [], [], [adjoint_region]
+            return [], [], [adjoint_region], dbg
 
         # Create a nested jaxpr scope for the body of the adjoint.
-        ctx = EvaluationContext.get_main_tracing_context()
-        with EvaluationContext.frame_tracing_context(ctx) as inner_trace:
+        with EvaluationContext.frame_tracing_context(debug_info=dbg) as inner_trace:
             in_classical_tracers, _ = tree_flatten((args, kwargs))
-            wffa, in_avals, _, _ = deduce_avals(self.target, args, kwargs)
+            wffa, in_avals, _, _ = deduce_avals(self.target, args, kwargs, debug_info=dbg)
             arg_classical_tracers = _input_type_to_tracers(inner_trace.new_arg, in_avals)
             with QueuingManager.stop_recording(), QuantumTape() as quantum_tape:
-                # FIXME: move all full_raise calls into a separate function
+                # FIXME: move all to_jaxpr_tracer calls into a separate function
                 res_classical_tracers = [
-                    inner_trace.full_raise(t)
+                    inner_trace.to_jaxpr_tracer(t)
                     for t in wffa.call_wrapped(*arg_classical_tracers)
                     if isinstance(t, DynamicJaxprTracer)
                 ]
@@ -442,7 +442,7 @@ def trace_body(self, args, kwargs):
                 inner_trace, quantum_tape, arg_classical_tracers, res_classical_tracers
             )
 
-        return in_classical_tracers, [], [adjoint_region]
+        return in_classical_tracers, [], [adjoint_region], wffa.debug_info
 
 
 class HybridAdjoint(HybridOp):
@@ -457,19 +457,20 @@ def trace_quantum(self, ctx, device, _trace, qrp) -> QRegPromise:
         body_trace = op.regions[0].trace
         body_tape = op.regions[0].quantum_tape
         res_classical_tracers = op.regions[0].res_classical_tracers
+        dbg = op.debug_info
 
         # Handle ops that were instantiated outside of a tracing context.
         if body_trace is None:
-            frame_ctx = EvaluationContext.frame_tracing_context(ctx)
+            frame_ctx = EvaluationContext.frame_tracing_context(debug_info=dbg)
         else:
-            frame_ctx = EvaluationContext.frame_tracing_context(ctx, body_trace)
+            frame_ctx = EvaluationContext.frame_tracing_context(body_trace)
 
         with frame_ctx as body_trace:
             qreg_in = _input_type_to_tracers(body_trace.new_arg, [AbstractQreg()])[0]
             qrp_out = trace_quantum_operations(body_tape, device, qreg_in, ctx, body_trace)
             qreg_out = qrp_out.actualize()
-            body_jaxpr, _, body_consts = ctx.frames[body_trace].to_jaxpr2(
-                res_classical_tracers + [qreg_out]
+            body_jaxpr, _, body_consts = body_trace.frame.to_jaxpr2(
+                res_classical_tracers + [qreg_out], dbg
             )
 
         qreg = qrp.actualize()
@@ -657,9 +658,9 @@ def ctrl_distribute(
 
     # Allow decompositions outside of a Catalyst context.
     if EvaluationContext.is_tracing():
-        ctx = EvaluationContext.get_main_tracing_context()
+        cur_trace = EvaluationContext.get_current_trace()
     else:
-        ctx = None
+        cur_trace = None
 
     new_ops = []
     for op in tape.operations:
@@ -675,8 +676,8 @@ def ctrl_distribute(
             else:
                 for region in [region for region in op.regions if region.quantum_tape is not None]:
                     # Re-enter a JAXPR frame but do not create a new one is none exists.
-                    if ctx and region.trace:
-                        trace_manager = EvaluationContext.frame_tracing_context(ctx, region.trace)
+                    if cur_trace and region.trace:
+                        trace_manager = EvaluationContext.frame_tracing_context(region.trace)
                     else:
                         trace_manager = nullcontext
 

frontend/catalyst/api_extensions/quantum_operators.py

@@ -152,7 +152,7 @@ def _decompose_nested_tapes(op, ctx, capabilities: DeviceCapabilities):
         if region.quantum_tape is None:
             new_tape = None
         else:
-            with EvaluationContext.frame_tracing_context(ctx, region.trace):
+            with EvaluationContext.frame_tracing_context(region.trace):
                 tapes, _ = catalyst_decompose(
                     region.quantum_tape, ctx=ctx, capabilities=capabilities
                 )

frontend/catalyst/device/decomposition.py

@@ -63,13 +63,12 @@ def _verify_nested(
 ) -> Any:
     """Traverse the nested quantum tape, carry a caller-defined state."""
 
-    ctx = EvaluationContext.get_main_tracing_context()
     for op in operations:
         inner_state = op_checker_fn(op, state)
         if has_nested_tapes(op):
             for region in nested_quantum_regions(op):
                 if region.trace is not None:
-                    with EvaluationContext.frame_tracing_context(ctx, region.trace):
+                    with EvaluationContext.frame_tracing_context(region.trace):
                         inner_state = _verify_nested(
                             region.quantum_tape.operations, inner_state, op_checker_fn
                         )