def state_init(inputs: jnp.ndarray,
               rng: jnp.ndarray,
               mask_fn: mask_utils.MaskFn,
               num_iter: int,
               keep_intermediates: bool = False):
  """Initializes the decoding state data structure."""
  cur_iter0 = jnp.array(0)
  token_dict = mask_fn(inputs, rng, jnp.ones((1,)), weight_mode='mask+refine')
  masked_tokens = token_dict['masked_inputs']
  uncond_masked_tokens = token_dict['uncond_masked_inputs']
  segment_ids = token_dict['segment_ids']
  unknown_mask = token_dict['weights'] == 1.
  if keep_intermediates:
    full_tokens0 = jnp.empty_like(
        inputs, shape=(inputs.shape[0], num_iter, *inputs.shape[1:]))
  else:
    full_tokens0 = jnp.empty_like(inputs[:, None])
  return State(
      cur_iter=cur_iter0,
      masked_tokens=masked_tokens,
      uncond_masked_tokens=uncond_masked_tokens,
      segment_ids=segment_ids,
      unknown_mask=unknown_mask,
      rng=rng,
      full_tokens=full_tokens0)
Initializes the decoding state data structure.
videogvt/train_lib/parallel_decode.py