code: def gather_topk_beams(
      nested, score_or_log_prob, batch_size, new_beam_size
  ):
      """Gathers the top-k beam slices given by score_or_log_prob array.
  
      Args:
        nested: pytree of arrays or scalars (the latter ignored).
        score_or_log_prob: [batch_size, old_beam_size] array of values to sort by
          for top-k selection of beam slices.
        batch_size: int: size of batch.
        new_beam_size: int: size of _new_ top-k selected beam dimension
  
      Returns:
        New pytree with new beam arrays containing top k new_beam_size slices.
        [batch_size, old_beam_size, ...] --> [batch_size, new_beam_size, ...]
      """
      _, topk_indices = lax.top_k(score_or_log_prob, k=new_beam_size)
      topk_indices = jnp.flip(topk_indices, axis=1)
      return gather_beams(
          nested, topk_indices, batch_size, new_beam_size
      )
  
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 7 for calling a global function with example usage: gather_topk_beams(nested, score_or_log_prob, batch_size, new_beam_size) and returns: gather_beams