PR #26262: [better_errors] Merge the JaxprDebugInfo and TracingDebugI…

…nfo into core.DebugInfo

Imported from GitHub PR jax-ml/jax#26262

Previously, we had two almost identical classes: `TracingDebugInfo` and `JaxprDebugInfo`. The only difference was that `TracingDebugInfo` had a thunk to return the result paths, while `JaxprDebugInfo` had the result paths resolved to a tuple. The separation of these types provided some clarity, but also led to code duplication and required conversions as the debugging info goes from `WrappedFun` to a `Jaxpr` and then to `WrappedFun` again.

Much of this PR is just renaming.
Copybara import of the project:

--
c70de6deedcc0311478936be232d042ede868d2f by George Necula <[email protected]>:

[better_errors] Merge the JaxprDebugInfo and TracingDebugInfo into core.DebugInfo

Previously, we had two almost identical classes: `TracingDebugInfo` and
`JaxprDebugInfo`. The only difference was that `TracingDebugInfo` had
a thunk to return the result paths, while `JaxprDebugInfo` had the
result paths resolved to a tuple. The separation of these types
provided some clarity, but also led to code duplication and
required conversions as the debugging info goes from `WrappedFun`
to a `Jaxpr` and then to `WrappedFun` again.

Merging this change closes #26262

PiperOrigin-RevId: 722684417

kfac_jax/_src/tag_graph_matcher.py

@@ -385,7 +385,7 @@ def make_jax_graph(
 
     debug_info = closed_jaxpr.jaxpr.debug_info
     if debug_info is not None:
-      debug_info = jax.core.JaxprDebugInfo(
+      debug_info = jax.core.DebugInfo(
           debug_info.traced_for,
           debug_info.func_src_info,
           debug_info.arg_names,