def loop_body_fn(state):
    """Beam search loop state update function."""
    rng = state.rng
    step = state.cur_iter
    prev_tokens = state.masked_tokens
    prev_uncond_masked_tokens = state.uncond_masked_tokens
    prev_segment_ids = state.segment_ids
    unknown_mask = state.unknown_mask

    # Calls model on current seqs to get next-iteration seqs.
    logits = decode_fn(prev_tokens, prev_segment_ids,
                       prev_uncond_masked_tokens)  # [batch, T, H, W, vocab]
    # Samples predicted tokens at current step.
    rng, sample_rng = jax.random.split(rng, 2)
    cur_tokens = sampling.sampling(
        logits,
        sample_rng,
        topk=sampling_topk,
        topp=sampling_topp,
        temperature=sampling_temperature)
    # Just updates the masked tokens.
    num_unknown = jnp.sum(unknown_mask, axis=(1, 2, 3))
    cur_tokens = jnp.where(unknown_mask, cur_tokens, prev_tokens)
    # Updates full_seqs with the current sampled_tokens.
    if keep_intermediates:
      full_tokens = state.full_seqs.at[:, step].set(cur_tokens)
    else:
      full_tokens = cur_tokens[:, None]

    # Computes the probabilities of each selected tokens.
    probs = jax.nn.softmax(logits, axis=-1)
    cur_probs = jnp.take_along_axis(probs, cur_tokens[..., None],
                                    -1)[..., 0]  # [batch, T, H, W]
    # Ignores the tokens given in the input by overwriting their confidence.
    cur_probs = jnp.where(unknown_mask, cur_probs, mask_utils.NEVER_MASK)
    # Mask ratio for the next round.
    progress = (step + 1) / num_iter
    mask_ratio = mask_schedule.schedule(progress, num_tokens,
                                        mask_scheduling_method)[None]

    # Adds noise to mask score for randomness
    rng, score_rng, mask_rng = jax.random.split(rng, 3)
    mask_score = jnp.log(cur_probs) + mask_temperature * (
        1. - progress) * jax.random.gumbel(score_rng, cur_probs.shape)
    # Mask tokens
    token_dict = mask_fn(
        cur_tokens,
        mask_rng,
        mask_ratio,
        weight_mode='mask+refine',
        mask_score=mask_score,
        num_masked_limits=(1, num_unknown - 1))
    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.
    return State(
        cur_iter=state.cur_iter + 1,
        masked_tokens=masked_tokens,
        uncond_masked_tokens=uncond_masked_tokens,
        segment_ids=segment_ids,
        unknown_mask=unknown_mask,
        rng=rng,
        full_tokens=full_tokens)
Beam search loop state update function.
videogvt/train_lib/parallel_decode.py