Python class

BatchProcessorInputs

`BatchProcessorInputs`

class max.interfaces.BatchProcessorInputs(logits, logit_offsets, context_batch)

source

Arguments for a batch logits processor.

logits: The model logits, a float32 tensor with shape (N_batch, vocab_size). N_batch is the number of logits returned by the model for each sequence in the batch.
logit_offsets: If the model returns multiple logits, this is a tensor with shape (batch_size + 1, 1) that contains the offsets of each sequence in the batch. Otherwise, this is None.
context_batch: The batch of contexts containing the inputs to the model.

Parameters:

context_batch: Sequence[TextGenerationContext]

source

The ordered sequence of generation contexts corresponding to each batch entry.

logit_offsets: md.Buffer | None

source

Offsets tensor with shape (batch_size + 1, 1) for multi-logit models, or None.

logits: md.Buffer

source

The model logits buffer with shape (N_batch, vocab_size) and float32 dtype.