For the complete documentation index, see llms.txt. Markdown versions of all pages are available by appending .md to any URL (e.g. /max/get-started.md).

Mojo function

compute_log_probabilities_ragged_shape

def compute_log_probabilities_ragged_shape[levels: Int](logits: ManagedTensorSlice[IOSpec[_, _].Input, static_spec=logits.static_spec], tokens: ManagedTensorSlice[IOSpec[_, _].Input, static_spec=tokens.static_spec], sampled_tokens: ManagedTensorSlice[IOSpec[_, _].Input, static_spec=sampled_tokens.static_spec], logit_row_offsets: ManagedTensorSlice[IOSpec[_, _].Input, static_spec=logit_row_offsets.static_spec], token_row_offsets: ManagedTensorSlice[IOSpec[_, _].Input, static_spec=token_row_offsets.static_spec], lp_output_offsets: ManagedTensorSlice[IOSpec[_, _].Input, static_spec=lp_output_offsets.static_spec], lp_output_offsets_host: ManagedTensorSlice[IOSpec[_, _].Input, static_spec=lp_output_offsets_host.static_spec]) -> IndexList[Int(2)]

Computes the output shapes for the ragged log-probabilities op.

Parameters:

levels (Int): Number of heap levels; the output second dimension is 2**levels.

Args:

logits (ManagedTensorSlice[IOSpec[_, _].Input, static_spec=logits.static_spec]): Input logits ragged by batch, shape [total_rows, vocab_size].
tokens (ManagedTensorSlice[IOSpec[_, _].Input, static_spec=tokens.static_spec]): Previously generated tokens across all batches, ragged.
sampled_tokens (ManagedTensorSlice[IOSpec[_, _].Input, static_spec=sampled_tokens.static_spec]): Most recently sampled token for each batch.
logit_row_offsets (ManagedTensorSlice[IOSpec[_, _].Input, static_spec=logit_row_offsets.static_spec]): Per-batch start offsets into the first axis of logits.
token_row_offsets (ManagedTensorSlice[IOSpec[_, _].Input, static_spec=token_row_offsets.static_spec]): Per-batch start offsets into tokens.
lp_output_offsets (ManagedTensorSlice[IOSpec[_, _].Input, static_spec=lp_output_offsets.static_spec]): Per-batch start offsets into the output row axis.
lp_output_offsets_host (ManagedTensorSlice[IOSpec[_, _].Input, static_spec=lp_output_offsets_host.static_spec]): Host-resident copy of lp_output_offsets whose last element gives the total number of output tokens.

Returns:

IndexList[Int(2)]: The output shape [num_output_tokens, 2**levels].