code: def noinline(
      fn: Callable, abstract_fn: Optional[Callable] = None
  ) -> Callable:  # pyright: ignore
      """Marks a function as not being inlined into a larger computation.
      This can help to reduce compile time at the expense of increased runtime.
  
      `fn` can be any callable `PyTree[Any] -> PyTree[Any]`. In addition `fn`
      itself may be any (callable) PyTree; any JAX arrays in `fn`'s PyTree are
      also treated as inputs to the noinline'd computation.
  
      `abstract_fn` determines the shapes/dtypes/pytrees of the output. (And must
       return results consistent with `fn`.) If `fn` is called as
      `fn(*args, **kwargs)`, then `abstract_fn` is called as
      ```python
      eqx.filter_eval_shape(
          abstract_fn, eqx.filter(fn, eqx.is_array), *args, **kwargs)
      )
      ```
      If not passed then `abstract_fn` is automatically constructed from `fn`.
      Specifying it is useful as noinline'd functions sharing the same
      `abstract_fn` may be substituted for each other without recompiling the
      enclosing computation graph; see the second example below.
  
      !!! Example
  
          ```python
          @noinline
          def f(x, y):
              print("Tracing!")
              return x + y
  
          @jax.jit
          def g(x, y):
              a = f(x, jnp.array(1))
              b = f(x, jnp.array(2))
              c = f(y, jnp.array(1))
              d = f(y, jnp.array(2))
              return a + b + c + d
  
          g(1, 1)
          ```
  
          In this example, `Tracing!` is printed twice. Once just for abstract
          evaluation (only) to figure out out its return shape in `f`'s
          computation graph. (Which is computationally cheap.) And then once
          when it is actually compiled as part of its own computation graph.
          (Which is computationally expensive.)
  
          Without the `noinline`, we would see `Tracing!` printed four times:
          once for each call, and with it being compiled each time. (Every time
          being computationally expensive.)
  
          Note how the `1` and the `2` are wrapped in arrays. If they were still
          Python scalars then they would be treated as static arguments, which
          will lead to recompilations of `f`. (E.g. what if there was an
          `if y == 1` command in there?)
  
      !!! Example
  
          ```python
          def abstract(_, x, y):
              return jnp.broadcast_arrays(x, y)[0]
  
          def f(x, y):
              print("Compiling f!")
              return x + y
  
          def g(x, y):
              print("Compiling g!")
              return x * y
  
          f = noinline(f, abstract)
          g = noinline(g, abstract)
  
          @jax.jit
          def call(fn, x, y):
              print("Compiling call!")
              return fn(x, y)
  
          call(f, 1, 1)  # Compiling call! Compiling f!
          call(g, 1, 1)  # Compiling g!
          ```
  
          In this example, we see how noinline'd functions with the same
          `abstract_fn` may be passed as inputs to the main call graph,
          and swapped without needing to recompile the main call graph.
  
      """
      dynamic_fn, static_fn = hashable_partition(fn, is_array)
      if abstract_fn is None:
  
          def abstract_fn(__dynamic_fn, *args, **kwargs):
              _fn = hashable_combine(__dynamic_fn, static_fn)
              return _fn(*args, **kwargs)
  
      try:
          dynamic_index = _fn_to_index[static_fn]
      except KeyError:
          dynamic_index = len(_index_to_fn)
          _fn_to_index[static_fn] = dynamic_index
          _index_to_fn.append(static_fn)
      dynamic_index = np.array(dynamic_index)
      noinline_fn = _NoInlineWrapper(
          dynamic_index, abstract_fn, dynamic_fn
      )
      return module_update_wrapper(noinline_fn)
  
imports: import functools
  from collections.abc import Callable
  from typing import Any, Optional, Union
  import jax
  import jax._src.pretty_printer
  import jax._src.source_info_util
  import jax.core
  import jax.interpreters.ad
  import jax.interpreters.batching
  import jax.interpreters.mlir
  import jax.numpy
  import jax.tree_util
  import numpy
  from jaxtyping import Array, Int, PyTree
  from .._ad import filter_jvp
  from .._caches import internal_lru_caches
  from .._compile_utils import hashable_combine, hashable_partition
  from .._eval_shape import filter_eval_shape
  from .._filters import combine, is_array, partition
  from .._module import field, Module, module_update_wrapper
  from .._vmap_pmap import filter_vmap
  from .None import _primitive
  from ._primitive import filter_primitive_batching, filter_primitive_bind, filter_primitive_def, filter_primitive_jvp, filter_primitive_transpose, materialise_zeros
label_desc: Use image node_id 22 for calling a global function with example usage: noinline(fn, abstract_fn) and returns: module_update_wrapper, _fn