code: def custom_vmap_batching(
      args_flat, dims, *, call, rule, in_tree, out_tree
  ):
      del call
      (axis_size,) = {
          x.shape[d] for x, d in zip(args_flat, dims) if d is not None
      }
      args_flat = map(maybe_bdim_at_front, args_flat, dims)
      flat_in_batched = [d is not not_mapped for d in dims]
  
      args = tree_unflatten(in_tree, args_flat)
      in_batched = tree_unflatten(in_tree, flat_in_batched)
      out, out_batched = call_rule(rule, axis_size, in_batched, args)
      flat_outs, tree1 = tree_flatten(out)
      flat_out_batched, tree2 = tree_flatten(out_batched)
      check_vmap_rule_trees(rule, out_tree, tree1, tree2)
      flat_out_dims = [0 if b else not_mapped for b in flat_out_batched]
      return flat_outs, flat_out_dims
  
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 8 for calling a global function with example usage: custom_vmap_batching(args_flat, dims) and returns: flat_outs,  flat_out_dims