code: def add_beam_dim(x, beam_size):
      """Creates new beam dimension in non-scalar array and tiles into it."""
      if x.ndim == 0:  # ignore scalars (e.g. cache index)
          return x
      x = jnp.expand_dims(x, axis=1)
      tile_dims = [1] * x.ndim
      tile_dims[1] = beam_size
      return jnp.tile(x, tile_dims)
  
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 2 for calling a global function with example usage: add_beam_dim(x, beam_size) and returns: jnp, x