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 class

LabelScorer

LabelScorer

class max.pipelines.sampling.LabelScorer(session, device, in_dtype, unpadded_vocab_size)

source

Bases: object

Gathers candidate-token log-probabilities from next-token logits.

Construction compiles the label-gather graph, which normalizes each logits row over the unpadded vocabulary and reads the log-probability at each candidate token id. It runs on the device the logits live on.

scorer = LabelScorer(session, device, DType.bfloat16, vocab_size)
scores = scorer.score(logits, [[1234, 5678], None])
# scores == [[logp_1234, logp_5678], None]

Parameters:

• session (InferenceSession) – The inference session used to load the graph.
• device (Device) – The device the model logits live on.
• in_dtype (DType) – The dtype of the model’s next-token logits.
• unpadded_vocab_size (int) – The tokenizer’s token count.

Raises:

ValueError – If unpadded_vocab_size is not positive.

score()

score(logits, label_token_ids)

source

Scores candidate labels for each row of logits.

Launches the graph once for the whole batch, then copies the gathered scores to the host, which blocks until the launch finishes. Returns without launching when no row has candidates.

Parameters:

• logits (Buffer) – [batch, vocab] next-token logits, one row per request.
• label_token_ids (Sequence[Sequence[int] | None]) – Per-row candidate token ids, or None for rows that are not scoring requests (those rows get None).

Returns:

Per-row full-vocabulary log-probabilities of the candidates, in candidate order, or None for non-scoring rows.

Raises:

ValueError – If logits has a different number of rows than label_token_ids, or a candidate id is outside the vocabulary.

Return type:

list[list[float] | None]