code: def flat_batch_beam_expand(x, beam_size):
      """Expands the each batch item by beam_size in batch_dimension."""
      return flatten_beam_dim(add_beam_dim(x, 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 5 for calling a global function with example usage: flat_batch_beam_expand(x, beam_size) and returns: flatten_beam_dim