code: def filter_make_jaxpr(
      fun: Callable[_P, Any],
  ) -> Callable[
      _P,
      tuple[
          jax.core.ClosedJaxpr,
          PyTree[jax.ShapeDtypeStruct],
          PyTree[Any],
      ],
  ]:
      """As `jax.make_jaxpr`, but accepts arbitrary PyTrees as input and output.
  
      **Arguments:**
  
      - `fun`: The function `fun(*arg, **kwargs)` whose jaxpr is to be computed. Its
          positional and keyword arguments may be anything, as can its return value.
  
      **Returns:**
  
      A wrapped version of `fun`, that when applied to example arguments
      `*args, **kwargs`, will return a 3-tuple of:
  
      - A `ClosedJaxpr` representing the evaluation of that function on those arguments.
      - A `PyTree[jax.ShapeDtypeStruct]` representing the output shape and dtype of the
          result.
      - A `PyTree[Any]` representing any non-array outputs from `fun`.
  
      The example arguments to be traced may be anything with `.shape` and `.dtype`
      fields (typically JAX arrays, NumPy arrays, of `jax.ShapeDtypeStruct`s). All
      other argments are treated statically. In particular, Python builtins (`bool`,
      `int`, `float`, `complex`) are treated as static inputs; wrap them in JAX/NumPy
      arrays if you would like them to be traced.
      """
      return module_update_wrapper(_MakeJaxpr(fun))
  
file_path: equinox/equinox/_make_jaxpr.py
imports: from collections.abc import Callable
  from typing import Any
  from typing_extensions import ParamSpec
  import jax
  import jax._src.traceback_util
  import jax.core
  import jax.tree_util
  from jaxtyping import PyTree
  from ._filters import combine, is_array, partition
  from ._module import Module, module_update_wrapper, Static
label_desc: Use image node_id 2 for calling a global function with example usage: filter_make_jaxpr(fun) and returns: module_update_wrapper