code: def _for_vmap(
      spmd_axis_name,
      axis_size,
      axis_name,
      main_type,
      args,
      dims,
      *,
      jaxpr,
      nsteps,
      reverse,
      which_linear,
      unroll,
  ):
      init_batched = [d is not batching.not_mapped for d in dims]
      closed_jaxpr = _cached_for_jaxpr(jaxpr)
      batched = init_batched
      for _ in range(len(batched)):
          _, out_batched = batching.batch_jaxpr(
              closed_jaxpr,
              axis_size,
              [False] + batched,
              instantiate=batched,
              axis_name=axis_name,
              spmd_axis_name=spmd_axis_name,
              main_type=main_type,
          )
          if out_batched == batched:
              break
          batched = map(operator.or_, batched, out_batched)
      else:
          raise Exception("Invalid fixpoint")
      args = [
          batching.broadcast(x, axis_size, 0)
          if now_bat and not was_bat
          else batching.moveaxis(x, d, 0)
          if now_bat
          else x
          for x, d, was_bat, now_bat in zip(
              args, dims, init_batched, batched
          )
      ]
      batched_jaxpr_, _ = batching.batch_jaxpr(
          pe.close_jaxpr(jaxpr),
          axis_size,
          [False] + batched,
          [],
          axis_name=axis_name,
          spmd_axis_name=spmd_axis_name,
          main_type=main_type,
      )
      batched_jaxpr, () = (
          batched_jaxpr_.jaxpr,
          batched_jaxpr_.consts,
      )  # TODO consts
      out_flat = for_p.bind(
          *args,
          jaxpr=batched_jaxpr,
          nsteps=nsteps,
          reverse=reverse,
          which_linear=which_linear,
          unroll=unroll,
      )
      return out_flat, [
          0 if b else batching.not_mapped for b in batched
      ]
  
imports: from __future__ import annotations
  from collections.abc import Sequence
  import functools
  import operator
  from typing import Any, Callable, Generic, TypeVar
  import jax.numpy
  from jax import lax
  from jax.api_util import flatten_fun_nokwargs
  from jax._src.interpreters import ad
  from jax._src.interpreters import batching
  from jax._src.interpreters import mlir
  from jax._src.interpreters import partial_eval
  from jax.tree_util import tree_flatten, tree_structure, tree_unflatten, treedef_tuple, tree_map, tree_leaves, PyTreeDef
  from jax._src import ad_util
  from jax._src import core
  from jax._src import dispatch
  from jax._src import dtypes
  from jax._src import linear_util
  from jax._src import source_info_util
  from jax._src.state.types import ReadEffect, AbstractRef, StateEffect
  from jax._src.state import discharge
  from jax._src.state import primitives
  from jax._src.state import utils
  from jax._src.state import types
  from jax._src.typing import Array
  from jax._src.util import partition_list, merge_lists, safe_map, safe_zip, split_list, split_dict, weakref_lru_cache
  from jax._src.lax.control_flow import loops
  from jax._src.lax.control_flow.common import _abstractify, _initial_style_jaxpr
label_desc: Use image node_id 10 for calling a global function with example usage: _for_vmap(spmd_axis_name, axis_size, axis_name, main_type, args, dims) and returns: out_flat,  unknown