[better_errors] Make it explicit that debug_info is not None.

Now all internal uses of lu.wrap_init and core.Jaxpr are with actual
debug info. This enables us to clean up the type declarations and
to remove the checks whether debug_info is present.

For usage outside of the JAX internals, we change
`jax.extend.linear_util.wrap_init` to be usable without debug_info,
for temporary backwards compatibility. We emit a deprecation
warning and fill-in some fake debugging info.

PiperOrigin-RevId: 725512692

Author: gnecula

jax/_src/linear_util.py

@@ -19,7 +19,8 @@
 
    from jax._src import linear_util as lu
 
-   wf = lu.wrap_init(f)  # Produce a WrappedFun for applying transformations on `f`
+   # Produce a WrappedFun for applying transformations on `f`
+   wf = lu.wrap_init(f, debug_info=api_util.debug_info("test", f, (), {}))
 
 A `WrappedFun` object represents a function `f`, together with a sequence of
 nested transformations that are to be applied to the positional and keyword
@@ -66,6 +67,7 @@ def trans1(static_arg, *dynamic_args, **kwargs):
 from collections.abc import Callable, Sequence
 from functools import partial
 from typing import Any, NamedTuple
+import warnings
 import weakref
 
 from jax._src import config
@@ -150,14 +152,15 @@ class WrappedFun:
     stores: a list of out_store for the auxiliary output of the `transforms`.
     params: extra parameters to pass as keyword arguments to `f`, along with the
       transformed keyword arguments.
+    debug_info: debugging info about the function being wrapped.
   """
   __slots__ = ("f", "f_transformed", "transforms", "stores", "params", "in_type", "debug_info")
 
   def __init__(self, f: Callable,
                f_transformed: Callable,
                transforms,
                stores: tuple[Store | EqualStore | None, ...], params, in_type,
-               debug_info: DebugInfo | None):
+               debug_info: DebugInfo):
     self.f = f
     self.f_transformed = f_transformed
     self.transforms = transforms
@@ -168,7 +171,7 @@ def __init__(self, f: Callable,
 
   @property
   def __name__(self):
-    return getattr(self.f, '__name__', '<unnamed wrapped function>')
+    return self.debug_info.func_name
 
   def wrap(self, gen, gen_static_args,
            out_store: Store | EqualStore | None) -> WrappedFun:
@@ -254,6 +257,7 @@ def fun_name(f):
   except:
     return str(f)
 
+
 class DebugInfo(NamedTuple):
   """Debugging info about a func, its arguments, and results."""
   traced_for: str             # e.g. 'jit', 'scan', etc
@@ -313,19 +317,27 @@ def filter_result_paths(self, keep: Sequence[bool]) -> tuple[str, ...]:
     return tuple(v for v, b in zip(self.safe_result_paths(len(keep)), keep) if b)
 
 
+def _missing_debug_info_msg():
+  warnings.warn(
+      "linear_util.wrap_init() or core.Jaxpr are called without a DebugInfo "
+      "object. This behavior is deprecated, use api_util.debug_info() to "
+      "construct a proper DebugInfo object.",
+      DeprecationWarning, stacklevel=2)
+  return DebugInfo("missing_debug_info", fun_name(f), (), {})
+
+
 def wrap_init(f: Callable, params=None, *,
-              debug_info: DebugInfo | None = None) -> WrappedFun:
+              debug_info: DebugInfo) -> WrappedFun:
   """Wraps function `f` as a `WrappedFun`, suitable for transformation."""
   params_dict = {} if params is None else params
   params = () if params is None else tuple(sorted(params.items()))
   fun = WrappedFun(f, partial(f, **params_dict), (), (), params, None, debug_info)
-  if debug_info:
-    if debug_info.result_paths is None:
-      fun, result_paths_thunk = _get_result_paths_thunk(fun)
-      debug_info = debug_info._replace(
-          result_paths=HashableFunction(result_paths_thunk, closure=()))
-    fun = WrappedFun(fun.f, fun.f_transformed, fun.transforms, fun.stores,
-                     fun.params, fun.in_type, debug_info)
+  if debug_info.result_paths is None:
+    fun, result_paths_thunk = _get_result_paths_thunk(fun)
+    debug_info = debug_info._replace(
+        result_paths=HashableFunction(result_paths_thunk, closure=()))
+  fun = WrappedFun(fun.f, fun.f_transformed, fun.transforms, fun.stores,
+                    fun.params, fun.in_type, debug_info)
   return fun
 
 

jax/extend/linear_util.py

@@ -15,6 +15,8 @@
 # Note: import <name> as <name> is required for names to be exported.
 # See PEP 484 & https://github.com/jax-ml/jax/issues/7570
 
+from typing import Callable
+
 from jax._src.linear_util import (
   StoreException as StoreException,
   WrappedFun as WrappedFun,
@@ -24,5 +26,17 @@
   transformation_with_aux as transformation_with_aux,
   transformation2 as transformation2,
   transformation_with_aux2 as transformation_with_aux2,
-  wrap_init as wrap_init,
+  wrap_init as _wrap_init,
+  _missing_debug_info_msg as _missing_debug_info_msg,
 )
+
+# Version of wrap_init that does not require a DebugInfo object.
+# This usage is deprecated, use api_util.debug_info() to construct a proper
+# DebugInfo object.
+def wrap_init(f: Callable, params=None, *, debug_info=None) -> WrappedFun:
+  return _wrap_init(
+      f, params,
+      debug_info=_missing_debug_info_msg() if debug_info is None else debug_info)
+
+del _wrap_init
+del _missing_debug_info_msg