code: def _noinline_batch(inputs, batch_axes):
      dynamic_index, abstract_fn, transforms, args = inputs
      (
          dynamic_index_bdim,
          abstract_fn_bdim,
          transforms_bdim,
          args_bdim,
      ) = batch_axes
      assert (
          len(jtu.tree_leaves((abstract_fn_bdim, transforms_bdim))) == 0
      )  # all none
      if dynamic_index_bdim is not batching.not_mapped:
          # The batch rule for `lax.cond` with vmap'd predicate simply
          # broadcasts all constants in the branches. In particular it may broadcast
          # this. We simply need to ignore this and return to having a single dynamic
          # index.
          # This is actually a silent error if you do something exceptionally silly, and
          # try to manually combine two different `noinline`d functions. There's no real
          # way to catch this that wouldn't slow things down at runtime though, I don't
          # think.
          assert jnp.ndim(dynamic_index) == 1
          dynamic_index = dynamic_index[0]
      del dynamic_index_bdim, abstract_fn_bdim, transforms_bdim
      args = jtu.tree_map(
          _move_to_front, args, args_bdim, is_leaf=_is_not_mapped
      )
      args_bdim = jtu.tree_map(
          _int_to_zero, args_bdim, is_leaf=_is_not_mapped
      )
      out = filter_primitive_bind(
          noinline_p,
          dynamic_index,
          abstract_fn,
          transforms + [_MetaBatchTransform(args_bdim)],
          args,
      )
      return out, jtu.tree_map(lambda _: 0, out)
  
imports: import functools
  from collections.abc import Callable
  from typing import Any, Optional, Union
  import jax
  import jax._src.pretty_printer
  import jax._src.source_info_util
  import jax.core
  import jax.interpreters.ad
  import jax.interpreters.batching
  import jax.interpreters.mlir
  import jax.numpy
  import jax.tree_util
  import numpy
  from jaxtyping import Array, Int, PyTree
  from .._ad import filter_jvp
  from .._caches import internal_lru_caches
  from .._compile_utils import hashable_combine, hashable_partition
  from .._eval_shape import filter_eval_shape
  from .._filters import combine, is_array, partition
  from .._module import field, Module, module_update_wrapper
  from .._vmap_pmap import filter_vmap
  from .None import _primitive
  from ._primitive import filter_primitive_batching, filter_primitive_bind, filter_primitive_def, filter_primitive_jvp, filter_primitive_transpose, materialise_zeros
label_desc: Use image node_id 18 for calling a global function with example usage: _noinline_batch(inputs, batch_axes) and returns: out,  jtu