code: def unflatten_beam_dim(x, batch_size, beam_size):
      """Unflattens the first, flat batch*beam dimension of a non-scalar array."""
      if x.ndim == 0:  # ignore scalars (e.g. cache index)
          return x
      assert batch_size * beam_size == x.shape[0]
      return x.reshape((batch_size, beam_size) + x.shape[1:])
  
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 4 for calling a global function with example usage: unflatten_beam_dim(x, batch_size, beam_size) and returns: x, x