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

Python function

logits_postprocess

logits_postprocess()

max.nn.transformer.logits_postprocess(h, input_row_offsets, return_n_logits, norm, lm_head, return_logits, return_hidden_states=ReturnHiddenStates.NONE, logits_scaling=1.0, capture_hidden_states=None)

source

Performs common logits postprocessing for single-device models.

Handles last-token gathering, logits computation, optional logits scaling, and hidden states return.

Parameters:

• h (TensorValue) – The hidden states from the final transformer layer.
• input_row_offsets (TensorValue) – The row offsets for ragged batching.
• return_n_logits (TensorValue) – The number of logits to return per sequence.
• norm (Callable[[TensorValue], TensorValue]) – The normalization function applied before the language model head.
• lm_head (Callable[[TensorValue], TensorValue]) – The language model head projection.
• return_logits (ReturnLogits) – Which logits to return.
• return_hidden_states (ReturnHiddenStates) – Which hidden states to return. Defaults to ReturnHiddenStates.NONE.
• logits_scaling (float) – A divisor applied to logits after projection. Logits are divided by this value before returning. Defaults to 1.0 (no scaling).
• capture_hidden_states (list[list[TensorValue]] | None) – For SELECTED_LAYERS mode, the per-layer captured hidden states gathered during the layer loop. See extract_hs().

Returns:

A tuple of (last_logits,), optionally extended with (logits, offsets) when returning multiple logits, and further extended with hidden states when requested. Hidden-state entries for ALL / ALL_NORMALIZED / SELECTED_LAYERS modes are returned as a list[TensorValue] (one entry per device); other entries are single TensorValue.

Return type:

tuple[TensorValue, …]