code: def gather_beams(nested, beam_indices, batch_size, new_beam_size):
      """Gathers the beam slices indexed by beam_indices into new beam array.
  
      Args:
        nested: pytree of arrays or scalars (the latter ignored).
        beam_indices: array of beam_indices
        batch_size: int: size of batch.
        new_beam_size: int: size of _new_ beam dimension.
  
      Returns:
        New pytree with new beam arrays.
        [batch_size, old_beam_size, ...] --> [batch_size, new_beam_size, ...]
      """
      batch_indices = jnp.reshape(
          jnp.arange(batch_size * new_beam_size) // new_beam_size,
          (batch_size, new_beam_size),
      )
  
      def gather_fn(x):
          if x.ndim == 0:  # ignore scalars (e.g. cache index)
              return x
          else:
              return x[batch_indices, beam_indices]
  
      return jax.tree_util.tree_map(gather_fn, nested)
  
file_path: scenic/scenic/projects/avatar/decode.py
imports: import typing
  import flax
  import jax
  from jax import lax
  import jax.numpy
  import numpy
label_desc: Use image node_id 6 for calling a global function with example usage: gather_beams(nested, beam_indices, batch_size, new_beam_size) and returns: jax, x, x