code: def vmap_unrestricted(
      f: lu.WrappedFun, *args, in_axes, axis_name, axis_size
  ):
      f, out_axes = batching.batch_subtrace(f)
      f = batching._batch_outer(
          f, axis_name, axis_size, in_axes, batching.BatchTrace, None
      )
      outs = f.call_wrapped(*args)
      return outs, out_axes()
  
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 6 for calling a global function with example usage: vmap_unrestricted(f) and returns: outs,  out_axes