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

DistributedLogitsPostprocessMixin

`DistributedLogitsPostprocessMixin`

class max.nn.transformer.DistributedLogitsPostprocessMixin

source

Mixin providing logits postprocessing for multi-device sharded models.

Requires: self.norm_shards, self.lm_head, self.return_logits, self.devices. Optional: self.return_hidden_states, self.logits_scaling.

devices: list[DeviceRef]

source

lm_head: Callable[[list[TensorValue], Sequence[BufferValue]], Sequence[TensorValue]]

source

logit_softcapping: float | None = None

source

logits_scaling: float = 1.0

source

norm_shards: Sequence[Callable[[TensorValue], TensorValue]]

source

return_hidden_states: ReturnHiddenStates = 'none'

source

return_logits: ReturnLogits

source