code: def vmap_module(module, in_axes=0, out_axes=0, num_batch_dims=1):
      """Vectorize a module.
  
      Args:
        module: the module to vectorize.
        in_axes: the `in_axes` argument passed to vmap. See `jax.vmap`.
        out_axes: the `out_axes` argument passed to vmap. See `jax.vmap`.
        num_batch_dims: the number of batch dimensions (how many times to apply vmap
          to the module).
  
      Returns:
        A vectorized module.
      """
      for _ in range(num_batch_dims):
          module = nn.vmap(
              module,
              variable_axes={"params": None},
              split_rngs={"params": False},
              in_axes=in_axes,
              out_axes=out_axes,
          )
  
      return module
  
file_path: nerfies/nerfies/model_utils.py
imports: from flax import linen
  from flax import optim
  from flax import struct
  from jax import lax
  from jax import random
  import jax.numpy
label_desc: Use image node_id 11 for calling a global function with example usage: vmap_module(module, in_axes, out_axes, num_batch_dims) and returns: module