code: def vmap(
      fn: Callable[..., Any],
      variable_axes: Mapping[CollectionFilter, InOutAxis],
      split_rngs: Mapping[PRNGSequenceFilter, bool],
      in_axes=0,
      out_axes=0,
      axis_size: Optional[int] = None,
      axis_name: Optional[str] = None,
      spmd_axis_name: Optional[str] = None,
      metadata_params: Dict[Any, Any] = {},
  ) -> Callable[..., Any]:
      """A lifted version of ``jax.vmap``.
  
      See ``jax.vmap`` for the unlifted batch transform in Jax.
  
      ``vmap`` can be used to add a batch axis to a scope function.
      For example we could create a version of ``dense`` with
      a batch axis that does not share parameters::
  
        batch_dense = lift.vmap(
            nn.dense,
            in_axes=(0, None),
            variable_axes={'params': 0},
            split_rngs={'params': True})
  
      By using ``variable_axes={'params': 0}``, we indicate that the
      parameters themselves are mapped over and therefore not shared along
      the mapped axis. Consequently, we also split the 'params' RNG,
      otherwise the parameters would be initialized identically along
      the mapped axis.
  
      Similarly, ``vmap`` could be use to add a batch axis with parameter
      sharing::
  
        batch_foo = lift.vmap(
            foo,
            in_axes=0, out_axes=0,
            variable_axes={'params': None},
            split_rngs={'params': False})
  
      Here we use ``variable_axes={'params': None}`` to indicate the parameter
      variables are shared along the mapped axis. Consequently, the 'params'
      RNG must also be shared.
  
      Args:
        fn: the function to be transformed.
        variable_axes: the variable collections that are lifted into the batching
          transformation. Use `None` to indicate a broadcasted collection or an
          integer to map over an axis.
        split_rngs: Split PRNG sequences will be different for each index of the
          batch dimension. Unsplit PRNGs will be broadcasted.
        in_axes: Specifies the mapping of the input arguments (see `jax.vmap).
        out_axes: Specifies the mapping of the return value (see `jax.vmap).
        axis_size: Specifies the size of the batch axis. This only needs to be
          specified if it cannot be derived from the input arguments.
        axis_name: Specifies a name for the batch axis. Can be used together with
          parallel reduction primitives (e.g. `jax.lax.pmean`, `jax.lax.ppermute`,
          etc.). Note, this is only used for pmap and shmap. For SPMD jit, you do
          not need to manually synchronize. Just make sure that the axes are
          correctly annotated and XLA:SPMD will insert the necessary collectives.
        spmd_axis_name: Axis name added to any pjit sharding constraints appearing
          in `fn`. See also
          https://github.com/google/flax/blob/main/flax/linen/partitioning.py.
        metadata_params: arguments dict passed to AxisMetadata instances in the
          variable tree.
  
      Returns:
        A vectorized version of the input scope function.
      """
      variable_in_axes, variable_out_axes = _split_in_out_axes(
          variable_axes
      )
      variable_in_groups, variable_in_axes = _unzip2(
          variable_in_axes.items()
      )
      variable_out_groups, variable_out_axes = _unzip2(
          variable_out_axes.items()
      )
      rng_groups, rng_splits = _unzip2(split_rngs.items())
      rng_axes = tuple(
          0 if rng_split else None for rng_split in rng_splits
      )
  
      def inner(
          scope_fn, repack_fn, variable_groups, rng_groups, *args
      ):
          def find_axis_size(axis, x):
              if axis is not None:
                  leaves = jax.tree_util.tree_leaves(x)
                  if leaves:
                      return leaves[0].shape[axis]
              return ()
  
          # split rngs
          axis_sizes = jax.tree_util.tree_map(
              find_axis_size,
              (variable_in_axes, in_axes),
              (variable_groups, args),
          )
          axis_sizes = set(jax.tree_util.tree_leaves(axis_sizes))
          if axis_size is None and len(axis_sizes) == 1:
              (d_axis_size,) = axis_sizes
          elif len(axis_sizes) > 1:
              raise ValueError(
                  f"Inconsistent batch axis sizes: {axis_sizes}"
              )
          elif axis_size is None:
              raise ValueError(
                  "axis_size should be specified manually."
              )
          else:
              d_axis_size = axis_size
          split_fn = lambda rng: random.split(rng, d_axis_size)
  
          rng_groups = tuple(
              tree_map_rngs(split_fn, rng_group) if split else rng_group
              for rng_group, split in zip(rng_groups, rng_splits)
          )
  
          new_variable_groups = []
          for var_group, axis in zip(variable_groups, variable_in_axes):
              if axis is not None:
                  new_variable_groups.append(
                      meta.remove_axis(var_group, axis, metadata_params)
                  )
              else:
                  new_variable_groups.append(var_group)
          variable_groups = tuple(new_variable_groups)
  
          @functools.partial(
              jax.vmap,
              in_axes=(variable_in_axes, rng_axes, in_axes),
              out_axes=(out_axes, variable_out_axes),
              axis_name=axis_name,
              axis_size=axis_size,
              spmd_axis_name=spmd_axis_name,
          )
          @functools.wraps(fn)
          def mapped(variable_groups, rng_groups, args):
              scope = scope_fn(variable_groups, rng_groups)
              y = fn(scope, *args)
              return y, repack_fn(scope)
  
          y, vars_out = mapped(variable_groups, rng_groups, args)
          new_vars_out = []
          for var_group, axis in zip(vars_out, variable_out_axes):
              if axis is not None:
                  new_vars_out.append(
                      meta.add_axis(var_group, axis, metadata_params)
                  )
              else:
                  new_vars_out.append(var_group)
          vars_out = tuple(new_vars_out)
          return y, vars_out
  
      return pack(
          inner,
          variable_in_groups,
          variable_out_groups,
          rng_groups,
          name="vmap",
      )
  
imports: import collections
  import dataclasses
  import functools
  from typing import Any, Callable, Dict, Generic, Iterable, List, Mapping, Optional, Sequence, Tuple, TypeVar, Union
  import warnings
  from flax import traceback_util
  import jax
  from jax import random
  from .None import axes_scan, meta
  from .frozen_dict import freeze, unfreeze
  from .scope import CollectionFilter, DenyList, Filter, PRNGSequenceFilter, Scope, group_collections, in_filter, intersect_filters, is_filter_empty, subtract_filters, union_filters
label_desc: Use image node_id 14 for calling a global function with example usage: vmap(fn, variable_axes, split_rngs, in_axes, out_axes, axis_size, axis_name, spmd_axis_name, metadata_params) and returns: pack, y,  vars_out, , y,  repack_fn, leaves