Description
Description
As part of improving JAX error messages we ensure that a core.DebugInfo object is passed to all calls to linear_util.wrap_init and to the core.Jaxpr constructor. We have made the necessary changes into the JAX core code base (e.g., #26078, #26313, #26348, and several other PRs).
There are still a number of other projects that invoke directly these functions (and other helper functions that now must take core.DebugInfo). Note that only JAX-internal functions take a core.DebugInfo; this change does not affect projects that use only public APIs.
For a limited time, a DeprecationWarning is printed if jax.extend.linear_util.wrap_init is used without debugging info.
To obtain a core.DebugInfo, one can call api_util.debug_info, or retrieve the debug_info fields of lu.WrappedFun or core.Jaxpr.
System info (python version, jaxlib version, accelerator, etc.)
N/A