code: def filter_pmap(
      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,
      donate: Literal["all", "warn", "none"] = "none",
      **pmapkwargs,
  ):
      """
      !!! warning
  
          JAX has now added more powerful parallelism APIs directly to the JIT interface.
          As such, using [`equinox.filter_jit`][] with sharded inputs is now recommended
          over `filter_pmap`. See also the
          [parallelism example](../../../examples/parallelism/).
  
      Parallelises a function. By default, all JAX/NumPy arrays are parallelised down
      their leading axis (i.e. axis index 0), and all other types are broadcast.
  
      `jax.pmap`, and thus `equinox.filter_pmap`, also compiles their function in the same
      way as `jax.jit`. By default, all JAX arrays are traced, and all other arguments are
      treated as static inputs.
  
      **Arguments:**
  
      For both `in_axes` and `out_axes`, then `int` indicates an array axis to parallelise
      over, `None` indicates that an argument should be broadcast (not parallelise
      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 parallelise. 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 parallelised 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 parallelise 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.
      - `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;
          - `'warn'`: as above, but don't suppress unused buffer warnings;
          - `'none'`: the default, disables buffer donation.
  
      **Returns:**
  
      The parallelised version of `fun`.
  
      !!! tip
  
          To parallelise 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_pmap
          def f(x, y):
              return x + y
  
          @eqx.filter_pmap(in_axes=(None, 1))
          def g(x, y):
              return x + y
  
          f(jnp.array([1, 2]), jnp.array([3, 4]))  # both args parallelised down axis 0
  
          f(jnp.array([1, 2]), 3)                  # first arg parallelised down axis 0
                                                   # second arg broadcasted (as it's not
                                                   # a JAX array)
  
          g(jnp.array(1), jnp.array([[2, 3]]))     # first arg broadcasted
                                                   # second arg parallelised down axis 1
          ```
      """
  
      if fun is sentinel:
          return ft.partial(
              filter_pmap,
              in_axes=in_axes,
              out_axes=out_axes,
              axis_name=axis_name,
              axis_size=axis_size,
              donate=donate,
              **pmapkwargs,
          )
  
      deprecated_0_10(pmapkwargs, "default")
      deprecated_0_10(pmapkwargs, "fn")
      deprecated_0_10(pmapkwargs, "args")
      deprecated_0_10(pmapkwargs, "kwargs")
      deprecated_0_10(pmapkwargs, "out")
      if any(
          x in pmapkwargs
          for x in ("static_broadcasted_argnums", "donate_argnums")
      ):
          raise ValueError(
              "`pmapkwargs` cannot contain either 'static_broadcasted_argnums' or "
              "'donate_argnums'"
          )
  
      if donate == "arrays":
          warnings.warn(
              "The `donate='arrays'` option to `filter_pmap` has been renamed to "
              "`donate='all'`",
              DeprecationWarning,
          )
          donate = "all"
      if donate not in {"all", "warn", "none"}:
          raise ValueError(
              "`filter_jit(..., donate=...)` must be one of 'all', 'warn', or 'none'"
          )
      filter_warning = True if donate == "all" else False
      if donate != "none":
          pmapkwargs["donate_argnums"] = (0,)
  
      pmap_wrapper = _PmapWrapper(
          _fun=fun,
          _in_axes=in_axes,
          _out_axes=out_axes,
          _axis_name=axis_name,
          _axis_size=axis_size,
          _filter_warning=filter_warning,
          _pmapkwargs=pmapkwargs,
      )
      return module_update_wrapper(pmap_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 16 for calling a global function with example usage: filter_pmap(fun) and returns: module_update_wrapper, ft