code: def filter_jit(
      fun=sentinel,
      *,
      donate: Literal[
          "all", "all-except-first", "warn", "warn-except-first", "none"
      ] = "none",
      **jitkwargs,
  ):
      """An easier-to-use version of `jax.jit`. All JAX and NumPy arrays are traced, and
      all other types are held static.
  
      **Arguments:**
  
      - `fun` is a pure function to JIT compile.
      - `donate` indicates whether the buffers of JAX arrays are donated or not. It
          should either be:
          - `'all'`: donate all arrays and suppress all warnings about unused buffers;
          - `'all-except-first'`: donate all arrays except for those in the first
              argument, and suppress all warnings about unused buffers;
          - `'warn'`: as above, but don't suppress unused buffer warnings;
          - `'warn-except-first'`: as above, but don't suppress unused buffer warnings;
          - `'none'`: no buffer donation. (This the default.)
  
      **Returns:**
  
      The JIT'd version of `fun`.
  
      !!! example
  
          ```python
          # Basic behaviour
          @eqx.filter_jit
          def f(x, y):  # both args traced if arrays, static if non-arrays
              return x + y, x - y
  
          f(jnp.array(1), jnp.array(2))  # both args traced
          f(jnp.array(1), 2)  # first arg traced, second arg static
          f(1, 2)  # both args static
          ```
  
      !!! info
  
          Donating arguments allows their underlying memory to be used in the
          computation. This can produce speed and memory improvements, but means that you
          cannot use any donated arguments again, as their underlying memory has been
          overwritten. (JAX will throw an error if you try to.)
  
      !!! info
  
          If you want to trace Python `bool`/`int`/`float`/`complex` as well then you
          can do this by wrapping them into a JAX array: `jnp.asarray(x)`.
  
          If you want to donate only some arguments then this can be done by setting
          `filter_jit(donate="all-except-first")` and then passing all arguments that you
          don't want to donate through the first argument. (Packing multiple values into
          a tuple if necessary.)
      """
  
      if fun is sentinel:
          return ft.partial(filter_jit, donate=donate, **jitkwargs)
  
      deprecated_0_10(jitkwargs, "default")
      deprecated_0_10(jitkwargs, "fn")
      deprecated_0_10(jitkwargs, "args")
      deprecated_0_10(jitkwargs, "kwargs")
      deprecated_0_10(jitkwargs, "out")
      if any(
          x in jitkwargs
          for x in (
              "static_argnums",
              "static_argnames",
              "donate_argnums",
          )
      ):
          raise ValueError(
              "`jitkwargs` cannot contain 'static_argnums', 'static_argnames' or "
              "'donate_argnums'"
          )
      signature = inspect.signature(fun)
  
      if donate == "all":
          filter_warning = True
          donate_first = True
          donate_rest = True
      elif donate == "all-except-first":
          filter_warning = True
          donate_first = False
          donate_rest = True
      elif donate == "warn":
          filter_warning = False
          donate_first = True
          donate_rest = True
      elif donate == "warn-except-first":
          filter_warning = False
          donate_first = False
          donate_rest = True
      elif donate == "none":
          filter_warning = False
          donate_first = False
          donate_rest = False
      else:
          raise ValueError(
              "`filter_jit(..., donate=...)` must be one of 'all', 'all-except-first', "
              "'warn', 'warn-except-first', or 'none'."
          )
  
      _, name = get_fn_names(fun)
      dynamic_fun, static_fun = hashable_partition(fun, is_array)
      cached = _filter_jit_cache(fun, jitkwargs)
  
      jit_wrapper = _JitWrapper(
          fn=name,
          _signature=signature,
          _dynamic_fun=dynamic_fun,
          _static_fun=static_fun,
          _cached=cached,
          filter_warning=filter_warning,
          donate_first=donate_first,
          donate_rest=donate_rest,
      )
      return module_update_wrapper(jit_wrapper)
  
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 8 for calling a global function with example usage: filter_jit(fun) and returns: module_update_wrapper, ft