Python module

log_probabilities

compute_log_probabilities_ragged()

max.pipelines.lib.log_probabilities.compute_log_probabilities_ragged(*, input_row_offsets, logits, next_token_logits, tokens, sampled_tokens, batch_top_n, batch_echo)

Computes the log probabilities for ragged model outputs.

Parameters:

• input_row_offsets (ndarray ) – Token offsets into token-indexed buffers, by batch index. Should have 1 more element than there are batches (batch n is token indices [input_row_offsets[n], input_row_offsets[n+1])).
• logits (ndarray | None ) – (tokens, vocab_dim) tensor full of tensor logits. Token dimension mapped to batches using input_row_offsets.
• next_token_logits (ndarray ) – (batch, vocab_dim) tensor full of logits for next tokens per batch.
• sampled_tokens (ndarray ) – (batch_dim,) tensor of sampled token per batch
• batch_top_n (Sequence [ int ] ) – Number of top log probabilities to return per input in the batch. For any element where top_n == 0, the LogProbabilities is skipped.
• batch_echo (Sequence [ bool ] ) – Whether to include input tokens in the returned log probabilities.
• tokens (ndarray )

Returns:

Computed log probabilities for each item in the batch.

Return type:

list[LogProbabilities | None]

log_softmax()

max.pipelines.lib.log_probabilities.log_softmax(x, axis=-1)

Compute the logarithm of the softmax function.

This implementation uses the identity log(softmax(x)) = x - log(sum(exp(x))) with numerical stability improvements to prevent overflow/underflow.

Parameters:

• x (ndarray ) – Input array
• axis (int ) – Axis to compute values along

Returns:

Array with same shape as x, representing log(softmax(x))

Return type:

ndarray

Was this page helpful?

Thank you! We'll create more content like this.

Thank you for helping us improve!

😔 What went wrong?

Some code doesn’t work

It includes inaccurate information

It's missing information I need

It was difficult to understand

Other

Submit