code: def filter_vmap(
      fun=sentinel,
      *,
      in_axes: PyTree[AxisSpec] = if_array(0),
      out_axes: PyTree[AxisSpec] = if_array(0),
      axis_name: Hashable = None,
      axis_size: Optional[int] = None,
      **vmapkwargs,
  ):
      """Vectorises a function. By default, all JAX/NumPy arrays are vectorised down their
      leading axis (i.e. axis index 0), and all other types are broadcast.
  
      **Arguments:**
  
      For both `in_axes` and `out_axes`, then `int` indicates an array axis to vectorise
      over, `None` indicates that an argument should be broadcast (not vectorised
      over), and callables `Leaf -> Union[None, int]` are mapped and evaluated on every
      leaf of their subtree. `None` should be used for non-JAX-array arguments.
  
      - `fun` is a pure function to vectorise. Should be of the form `fun(*args)`; that
          is to say it cannot accept keyword arguments.
      - `in_axes` indicates which axes of the input arrays should be vectorised over.
          It should be a PyTree of `None`, `int`, or callables `Leaf -> Union[None, int]`.
          Its tree structure should either be:
          1. a prefix of the input tuple of `args`.
          2. a dictionary, in which case the named argments will use the specified indices
              to vectorise over, and all other arguments will have the default
              `eqx.if_array(0)`.
      - `out_axes` indicates which axis of the output arrays the mapped axis should appear
          at. It should be a PyTree of `None`, `int`, or callables
          `Leaf -> Union[None, int]`, and its tree structure should be a prefix of the
          output `fun(*args)`.
      - `axis_name` is an optional hashable Python object used to identify the mapped
          axis so that parallel collectives (e.g. `jax.lax.psum`) can be applied.
      - `axis_size` is an optional `int` describing the size of the axis mapped. This
          only needs to be passed if none of the input arguments are vectorised, as else
          it can be deduced by looking at the argument shapes.
  
      **Returns:**
  
      The vectorised version of `fun`.
  
      !!! tip
  
          To vectorise all JAX/NumPy arrays down their `j`th axis, and broadcast all other
          types, then you can use `equinox.if_array(j)`, which returns a callable
          `leaf -> j if is_array(leaf) else None`. For example: the default values of
          `in_axes` and `out_axes` are both `equinox.if_array(0)`.
  
      !!! example
  
          ```python
          import equinox as eqx
          import jax.numpy as jnp
  
          @eqx.filter_vmap
          def f(x, y):
              return x + y
  
          @eqx.filter_vmap(in_axes=(None, 1))
          def g(x, y):
              return x + y
  
          f(jnp.array([1, 2]), jnp.array([3, 4]))  # both args vectorised down axis 0
  
          f(jnp.array([1, 2]), 3)                  # first arg vectorised down axis 0
                                                   # second arg broadcasted
  
          g(jnp.array(1), jnp.array([[2, 3]]))     # first arg broadcasted
                                                   # second arg vectorised down axis 1
          ```
  
      !!! example
  
          `filter_vmap` can be used to easily create ensembles of models. For example,
          here's an ensemble of eight MLPs:
  
          ```python
          import equinox as eqx
          import jax.random as jr
  
          key = jr.PRNGKey(0)
          keys = jr.split(key, 8)
  
          # Create an ensemble of models
  
          @eqx.filter_vmap
          def make_ensemble(key):
              return eqx.nn.MLP(2, 2, 2, 2, key=key)
  
          mlp_ensemble = make_ensemble(keys)
  
          # Evaluate each member of the ensemble on the same data
  
          @eqx.filter_vmap(in_axes=(eqx.if_array(0), None))
          def evaluate_ensemble(model, x):
              return model(x)
  
          evaluate_ensemble(mlp_ensemble, jr.normal(key, (2,)))
  
          # Evaluate each member of the ensemble on different data
  
          @eqx.filter_vmap
          def evaluate_per_ensemble(model, x):
              return model(x)
  
          evaluate_per_ensemble(mlp_ensemble, jr.normal(key, (8, 2)))
          ```
  
          Here, `make_ensemble` works because [`equinox.nn.MLP`][] is a PyTree, and so it
          is a valid output from a `filter_vmap`. This PyTree includes some JAX arrays
          (the weights and biases) and some non-JAX-arrays (e.g. activation functions).
          `filter_vmap` will vectorise the JAX arrays (with separate weights for each
          member of the ensemble) whilst leaving the non-JAX-arrays alone.
  
          Note that as the weights in `mlp_ensemble` now have a leading batch dimension
          -- that the weights of `eqx.nn.MLP` instances do not typically have -- then it
          cannot be called directly. It must instead be passed back into a vectorised
          region to be called.
      """
  
      if fun is sentinel:
          return ft.partial(
              filter_vmap,
              in_axes=in_axes,
              out_axes=out_axes,
              axis_name=axis_name,
              axis_size=axis_size,
              **vmapkwargs,
          )
  
      deprecated_0_10(vmapkwargs, "default")
      deprecated_0_10(vmapkwargs, "fn")
      deprecated_0_10(vmapkwargs, "args")
      deprecated_0_10(vmapkwargs, "kwargs")
      deprecated_0_10(vmapkwargs, "out")
  
      vmap_wrapper = _VmapWrapper(
          _fun=fun,
          _in_axes=in_axes,
          _out_axes=out_axes,
          _axis_name=axis_name,
          _axis_size=axis_size,
          _vmapkwargs=vmapkwargs,
      )
      return module_update_wrapper(vmap_wrapper)
  
imports: import dataclasses
  import functools
  import inspect
  import warnings
  from collections.abc import Callable, Hashable
  from typing import Any, Literal, Optional, overload, Union
  import jax
  import jax._src.traceback_util
  import jax.core
  import jax.interpreters.batching
  import jax.interpreters.pxla
  import jax.numpy
  import jax.tree_util
  import numpy
  from jaxtyping import PyTree
  from ._compile_utils import compile_cache, hashable_combine, 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, filter, is_array, is_array_like, partition
  from ._module import Module, module_update_wrapper, Partial, Static
label_desc: Use image node_id 9 for calling a global function with example usage: filter_vmap(fun) and returns: module_update_wrapper, ft