code: def _call(self, is_lower, args, kwargs):
      maybe_dummy = _common_preprocess(self._axis_size, kwargs)
      del kwargs
  
      in_axes = _named_in_axes(self._fun, self._in_axes, args)
      in_axes = _resolve_axes(args, in_axes)
      in_axes = (None, in_axes, 0, None)
  
      dynamic, static = partition(
          (self._fun, args, maybe_dummy, self._out_axes), is_array
      )
      struct = jtu.tree_map(
          lambda x: jax.ShapeDtypeStruct(x.shape, x.dtype), dynamic
      )
  
      cached = _filter_pmap_cache(
          self._fun,
          static,
          struct,
          in_axes,
          self._axis_name,
          self._axis_size,
          self._pmapkwargs,
      )
  
      if is_lower:
          return Lowered(
              cached.lower(dynamic),
              (self._fun, self._out_axes, self._axis_size),
              _preprocess,  # pyright: ignore
              _postprocess,  # pyright: ignore
          )
      else:
          if self._filter_warning is True:
              with warnings.catch_warnings():
                  warnings.filterwarnings(
                      "ignore",
                      message="Some donated buffers were not usable*",
                  )
                  out = cached(dynamic)
          else:
              out = cached(dynamic)
          return _postprocess(out)
  
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 _PmapWrapper obj's underlying member method code with example usage: obj._call(is_lower, args, kwargs) and returns: Lowered, _postprocess