code: def filter_vmap(
      fun: Callable[..., Any],
      *,
      in_axes: PyTree[AxisSpec] = if_array(0),
      out_axes: PyTree[AxisSpec] = if_array(0),
      axis_name: Hashable = None,
      axis_size: Optional[int] = None,
  ) -> Callable[..., Any]:
      ...
  
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 8 for calling a global function with example usage: filter_vmap(fun) without return types