code: def flatten_beam_dim(x):
      """Flattens the first two dimensions of a non-scalar array."""
      if x.ndim == 0:  # ignore scalars (e.g. cache index)
          return x
      return x.reshape((x.shape[0] * x.shape[1],) + x.shape[2:])
  
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 3 for calling a global function with example usage: flatten_beam_dim(x) and returns: x, x