code: def sequential_vmap(f):
      f = custom_vmap(f)
  
      @f.def_vmap
      def rule(axis_size, in_batched, *args):
          del axis_size
  
          def to_map(mapped_args):
              args = tree_merge(in_batched, mapped_args, bcast_args)
              return f(*args)
  
          mapped_args, bcast_args = tree_split(in_batched, list(args))
          out = lax.map(to_map, mapped_args)
          out_batched = tree_map(lambda _: True, out)
          return out, out_batched
  
      return f
  
imports: from __future__ import annotations
  import functools
  import operator
  from typing import Callable
  from jax import lax
  from jax._src import api
  from jax._src import core
  from jax._src import custom_api_util
  from jax._src import linear_util
  from jax._src import source_info_util
  from jax._src import traceback_util
  from jax._src import tree_util
  from jax._src import util
  from jax._src.api_util import flatten_fun_nokwargs
  from jax._src.interpreters import ad
  from jax._src.interpreters import batching
  from jax._src.interpreters.batching import not_mapped
  from jax._src.interpreters import mlir
  from jax._src.interpreters import partial_eval
  from jax._src.interpreters import xla
  from jax._src.tree_util import tree_flatten, tree_map, tree_structure, tree_unflatten, treedef_tuple
label_desc: Use image node_id 13 for calling a global function with example usage: sequential_vmap(f) and returns: f, out,  out_batched, f