Function sample_decode_loop_body_fn

axlearn/common/decoding.py:1055–1132

Sample decode loop state update function.

(state: DecodingState)

        return jnp.any(not_at_end & (~terminate_early))

    def sample_decode_loop_body_fn(state: DecodingState) -> DecodingState:
        """Sample decode loop state update function."""
        # [batch].
        cur_index = state.cur_index

        # Flatten the num_decodes dimension for the cache and the input IDs for this step.
        # [batch * num_decodes, 1].
        flat_ids = flatten_decoding_dim(
            jnp.take_along_axis(state.sequences, cur_index[:, None, None], axis=2)
        )
        # {[batch, num_decodes, ...], ...} --> {[batch * num_decodes, ...], ...}.
        flat_cache = vectorized_tree_map(flatten_decoding_dim, state.cache)

        # Call model on current tokens to get next-position logits and then unflatten.
        new_flat_log_probs, updated_flat_cache = tokens_to_scores(flat_ids, flat_cache)
        # [batch * num_decodes, vocab] --> [batch, num_decodes, vocab].
        candidate_log_probs = unflatten_decoding_dim(new_flat_log_probs, batch_size, num_decodes)
        # {[batch * num_decodes, ...], ...} --> {[batch, num_decodes, ...], ...}.
        updated_cache = vectorized_tree_map(
            lambda x: unflatten_decoding_dim(x, batch_size, num_decodes), updated_flat_cache
        )

        # Sample next token IDs according to logits.
        prng_key, updated_prng_key = jax.random.split(state.prng_key)
        # [batch, num_decodes].
        next_token = jax.random.categorical(prng_key, logits=candidate_log_probs, axis=2)

        # We allow next_index to exceed `max_decode_len-1`:
        # - When reading from next_index, mode="clip" will effectively read `max_decode_len-1`;
        # - When writing to next_index, one-hot will cause the write to become a no-op.
        # [batch].
        next_index = cur_index + 1

        # [batch, num_decodes].
        next_token_log_prob = jnp.sum(
            candidate_log_probs * jax.nn.one_hot(next_token, candidate_log_probs.shape[-1]),
            axis=-1,
        )

        # If end sequence tokens already emitted, adjust next token to pad_id and update log_prob.
        next_token = (
            next_token * ~state.stop_decoding
            + jnp.full_like(next_token, pad_id) * state.stop_decoding
        )
        # [batch, num_decodes].
        next_token_log_prob = (
            next_token_log_prob * ~state.stop_decoding
            + jnp.zeros_like(next_token_log_prob) * state.stop_decoding
        )

        # Update score and sequence trackers. For indices in `next_index` that exceed
        # `max_decode_len-1`, one-hot will zero-out the update.
        # [batch, num_decodes=1, length].
        oh_indices = jax.nn.one_hot(
            next_index[:, None], state.sequences.shape[-1], dtype=state.sequences.dtype
        )
        # [batch, num_decodes, length].
        updated_sequences = (