def state_init(inputs: jnp.ndarray,
               cache: Dict[str, jnp.array],
               rng: jnp.ndarray):
  """Initializes the decoding state data structure."""
  # assume each example has the same conditional lengths
  step = jnp.array(0)
  outputs = jnp.zeros((inputs.shape[0], inputs.shape[1] + 1), inputs.dtype)
  outputs = outputs.at[:, :inputs.shape[1]].set(inputs)
  score = jnp.zeros((inputs.shape[0],), jnp.float32)
  return State(
      step=step,
      rng=rng,
      outputs=outputs,
      cache=cache,
      score=score)
Initializes the decoding state data structure.
videogvt/train_lib/lm_sample_decode.py