code: def jaxpr_debug_info(
      jaxpr: core.Jaxpr,
      trace_debug: TracingDebugInfo | None,
      result_paths: tuple[str | None, ...] | None = None,
  ) -> core.Jaxpr:
      """Add debug info to jaxpr, given trace-time debug info and result paths."""
      if trace_debug is None:
          return jaxpr
      assert (result_paths is not None) ^ (
          trace_debug.result_paths is not None
      )
      if result_paths is None:
          result_paths = trace_debug.result_paths()  # type: ignore
      debug_info = core.JaxprDebugInfo(
          trace_debug.traced_for,
          trace_debug.func_src_info,
          trace_debug.arg_names,
          tuple(result_paths),
      )
      return jaxpr.replace(debug_info=debug_info)
  
imports: from __future__ import annotations
  from collections.abc import Iterable, Sequence
  import inspect
  import operator
  from functools import partial
  from typing import Any, Callable
  import warnings
  import numpy
  from jax._src import core
  from jax._src import dtypes
  from jax._src.abstract_arrays import numpy_scalar_types
  from jax._src.core import ShapedArray
  from jax._src.tree_util import PyTreeDef, tree_flatten, tree_unflatten, tree_map, tree_structure, treedef_children, generate_key_paths, keystr, broadcast_prefix, prefix_errors
  from jax._src.tree_util import _replace_nones
  from jax._src import linear_util
  from jax._src.linear_util import TracingDebugInfo
  from jax._src.util import safe_map, WrapKwArgs, Hashable, HashableFunction, Unhashable
  from jax._src import traceback_util
label_desc: Use image node_id 40 for calling a global function with example usage: jaxpr_debug_info(jaxpr, trace_debug, result_paths) and returns: jaxpr, jaxpr