code: def nfold_vmap(fun, N, *, broadcasted=True, in_axes=0):
      """Convenience function to apply (broadcasted) vmap N times."""
      _vmap = broadcasting_vmap if broadcasted else vmap
      for _ in range(N):
          fun = _vmap(fun, in_axes=in_axes)
      return fun
  
imports: import functools
  from typing import Any, NamedTuple, Union
  import numpy
  import jax
  from jax import lax
  from jax import tree_util
  from jax import vmap
  from jax._src import core
  from jax._src import dtypes
  from jax._src import stages
  from jax._src.api_util import flatten_axes
  import jax.numpy
  from jax.util import safe_zip
  from jax._src.lax.lax import _dot_general_shape_rule, DotDimensionNumbers
  from jax._src.typing import Array
label_desc: Use image node_id 1 for calling a global function with example usage: nfold_vmap(fun, N) and returns: fun