code: def beam_init(batch_size, beam_size, max_decode_len, cache):
      """Initializes the beam search state data structure."""
      cur_index0 = jnp.array(0)
      live_logprobs0 = jnp.tile(
          jnp.array([0.0] + [NEG_INF] * (beam_size - 1)),
          [batch_size, 1],
      )
      finished_scores0 = jnp.ones((batch_size, beam_size)) * NEG_INF
      live_seqs0 = jnp.zeros(
          (batch_size, beam_size, max_decode_len), jnp.int32
      )
      finished_seqs0 = jnp.zeros(
          (batch_size, beam_size, max_decode_len), jnp.int32
      )
      finished_flags0 = jnp.zeros((batch_size, beam_size), jnp.bool_)
      # add beam dimension to attention cache pytree elements
      beam_cache0 = jax.tree_util.tree_map(
          lambda x: add_beam_dim(x, beam_size), cache
      )
      return BeamState(
          cur_index=cur_index0,
          live_logprobs=live_logprobs0,
          finished_scores=finished_scores0,
          live_seqs=live_seqs0,
          finished_seqs=finished_seqs0,
          finished_flags=finished_flags0,
          cache=beam_cache0,
      )
  
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 8 for calling a global function with example usage: beam_init(batch_size, beam_size, max_decode_len, cache) and returns: BeamState