code: def _modify_traceback(e: Exception):
      # Remove JAX's UnfilteredStackTrace, with its huge error messages.
      e.__cause__ = None
      # Remove _JitWrapper.__call__ and _JitWrapper._call and Method.__call__ from the
      # traceback
      tb = (
          e.__traceback__
      ) = e.__traceback__.tb_next.tb_next.tb_next  # pyright: ignore
      try:
          # See https://github.com/google/jax/blob/69cd3ebe99ce12a9f22e50009c00803a095737c7/jax/_src/traceback_util.py#L190  # noqa: E501
          jax.lib.xla_extension.replace_thread_exc_traceback(
              tb
          )  # pyright: ignore
      except AttributeError:
          pass
      # IPython ignores __tracebackhide__ directives for the frame that actually raises
      # the error. We fix that here.
      try:
          get_ipython()  # pyright: ignore
      except NameError:
          pass
      else:
          import IPython  # pyright: ignore
  
          # Check that IPython supports __tracebackhide__
          if IPython.version_info[:2] >= (7, 17):  # pyright: ignore
              tb_stack = []
              while tb is not None:
                  tb_stack.append(tb)
                  tb = tb.tb_next
              for tb in reversed(tb_stack):
                  if not tb.tb_frame.f_locals.get(
                      "__tracebackhide__", False
                  ):
                      tb.tb_next = None
                      break
              else:
                  e.__traceback__ = None
  
imports: import functools
  import inspect
  import warnings
  from collections.abc import Callable
  from typing import Any, Literal, overload, TypeVar
  from typing_extensions import ParamSpec
  import jax
  import jax._src.traceback_util
  import jax.core
  from jaxtyping import PyTree
  from ._compile_utils import compile_cache, get_fn_names, hashable_combine, hashable_filter, hashable_partition, Lowered
  from ._custom_types import sentinel
  from ._deprecate import deprecated_0_10
  from ._doc_utils import doc_remove_args
  from ._filters import combine, is_array, partition
  from ._module import field, Module, module_update_wrapper, Partial, Static
label_desc: Use image node_id 5 for calling a global function with example usage: _modify_traceback(e) without return types