code: def _filter_pmap_cache(
      fun_names,
      static,
      struct,
      in_axes,
      axis_name,
      axis_size,
      pmapkwargs,
  ):
      def fun_abstract(_dynamic):
          fun, args, _, _ = combine(_dynamic, static)
          out = fun(*args)
          return filter(out, is_array_like)
  
      fun_abstract = jax.vmap(
          fun_abstract,
          in_axes=(in_axes,),
          axis_name=axis_name,
          axis_size=axis_size,
      )
      struct_out = jax.eval_shape(fun_abstract, struct)
      max_out_size = jtu.tree_reduce(
          lambda x, y: max(x, y.ndim), struct_out, 0
      )
      del fun_abstract, struct, struct_out
  
      def _check_map_out_axis(x: Optional[int]):
          if isinstance(x, int):
              if x < -max_out_size or x >= max_out_size:
                  raise ValueError(
                      "integers in filter_pmap(..., out_axes=...) must correspond to a "
                      "dimension of the output array"
                  )
          elif x is not None:
              raise ValueError(
                  "filter_pmap(..., out_axes=...) must contain only integers and Nones"
              )
  
      def fun_wrapped(_dynamic):
          _main = jax.core.find_top_trace(jtu.tree_leaves(_dynamic))
          _fun, _args, _, _out_axes = combine(_dynamic, static)
          _out = _fun(*_args)
          _out_axes = _bind_main(_main, _out_axes)
          _out_axes = _resolve_axes(_out, _out_axes)
          jtu.tree_map(_check_map_out_axis, _out_axes)
          _pmapd = []
          for i in range(-max_out_size, max_out_size):
              _i_axes = jtu.tree_map(lambda a: a == i, _out_axes)
              _pmapd.append(filter(_out, _i_axes))
          _none_axes = jtu.tree_map(
              _is_none, _out_axes, is_leaf=_is_none
          )
          _nonpmapd = filter(_out, _none_axes)
          _dynamic_nonpmapd, _static_nonpmapd = hashable_partition(
              _nonpmapd, is_array
          )
          return _pmapd, _dynamic_nonpmapd, Static(_static_nonpmapd)
  
      fun_name, fun_qualname = fun_names
      fun_wrapped.__name__ = fun_name
      fun_wrapped.__qualname__ = fun_qualname
  
      out_axes = (list(range(-max_out_size, max_out_size)), None, None)
  
      return jax.pmap(
          fun_wrapped,
          in_axes=(in_axes,),  # pyright: ignore
          out_axes=out_axes,  # pyright: ignore
          axis_name=axis_name,
          axis_size=axis_size,
          **pmapkwargs,
      )
  
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 10 for calling a global function with example usage: _filter_pmap_cache(fun_names, static, struct, in_axes, axis_name, axis_size, pmapkwargs) and returns: jax, filter, _pmapd,  _dynamic_nonpmapd,  Static