code: def __call__(self, /, *args, **kwargs):
      if len(kwargs) != 0:
          raise RuntimeError(
              "keyword arguments cannot be used with functions wrapped with "
              "`filter_vmap`"
          )
      del kwargs
  
      in_axes = _named_in_axes(self._fun, self._in_axes, args)
  
      # JAX only actually supports passing array-typed values as inputs.
      # Interestingly it's usually still fine to pass non-array-typed values as
      # broadcast inputs. *Unless* you mess up with an inconsistent axis size, in
      # which case it tries to get the axis sizes of all inputs and then fails on the
      # non-array inputs.
      # e.g.:
      # ```
      # jax.vmap(lambda x, y, z: 0,
      #          in_axes=(0, 0, None))(jnp.arange(3), jnp.arange(5), object())
      # ```
      in_axes = _resolve_axes(args, in_axes)
      unmapped_axis = jtu.tree_map(_is_none, in_axes, is_leaf=_is_none)
      static_args, dynamic_args = partition(args, unmapped_axis)
  
      def _fun_wrapper(_dynamic_args):
          _main = jax.core.find_top_trace(
              jtu.tree_leaves(_dynamic_args)
          ).main
          _args = combine(_dynamic_args, static_args)
          _out = self._fun(*_args)
          _out_axes = _bind_main(_main, self._out_axes)
          _out_axes = _resolve_axes(_out, _out_axes)
          _none_axes = jtu.tree_map(
              _is_none, _out_axes, is_leaf=_is_none
          )
          _nonvmapd, _vmapd = partition(_out, _none_axes)
          return _vmapd, Static((_nonvmapd, _out_axes))
  
      if len(jtu.tree_leaves(in_axes)) == 0 and self._axis_size is None:
          vmapd, static = _fun_wrapper(dynamic_args)
          if len(jtu.tree_leaves(vmapd)) != 0:
              raise ValueError(
                  "Cannot resolve batch dimension. Non-`None` `out_axes` requires "
                  "either `in_axes` or `axis_size` to be not `None`."
              )
      else:
          vmapd, static = jax.vmap(
              _fun_wrapper,
              in_axes=(in_axes,),
              out_axes=(0, None),
              axis_name=self._axis_name,
              axis_size=self._axis_size,
              **self._vmapkwargs,
          )(dynamic_args)
      nonvmapd, out_axes = static.value
  
      assert jtu.tree_structure(vmapd) == jtu.tree_structure(out_axes)
      vmapd = jtu.tree_map(_swapaxes, vmapd, out_axes)
  
      return combine(vmapd, nonvmapd)
  
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 2 for calling the _VmapWrapper obj's underlying member method code with example usage: obj.__call__() and returns: combine, _vmapd,  Static