def loop_body_fn(state):
    """Beam search loop state update function."""
    rng = state.rng
    step = state.step
    cache = state.cache
    outputs = state.outputs

    logits, new_cache = decode_fn(
        outputs[:, step, None],
        cache,
        step >= prefix_lengths - 1)
    # Samples predicted tokens at current step.
    new_rng, sample_rng = jax.random.split(rng, 2)
    cur_tokens = sampling.sampling(
        logits,
        sample_rng,
        topk=sampling_topk,
        topp=sampling_topp,
        temperature=sampling_temperature)
    new_ids = jnp.where(step < prefix_lengths - 1,
                        outputs[:, step+1, None],
                        cur_tokens)
    log_probs = jax.nn.log_softmax(logits, axis=-1)
    cur_tokens_log_probs = jnp.take_along_axis(
        jnp.reshape(log_probs, (log_probs.shape[0], -1)),
        cur_tokens,
        axis=-1)
    cur_log_probs = jnp.where(
        step < prefix_lengths - 1,
        jnp.ones((log_probs.shape[0], 1), log_probs.dtype),
        cur_tokens_log_probs)
    new_outputs = outputs.at[:, step+1].set(new_ids[:, 0])
    new_score = state.score + jnp.squeeze(cur_log_probs, -1)
    return State(
        step=step+1,
        rng=new_rng,
        outputs=new_outputs,
        cache=new_cache,
        score=new_score)
Beam search loop state update function.
videogvt/train_lib/lm_sample_decode.py