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

ragged_batch_inputs

ragged_batch_inputs()

max.pipelines.lib.interfaces.ragged_batch_inputs(device_ref, fields=('tokens', 'row_offsets', 'return_n_logits'), *, data_parallel_degree=1, hidden_states=None, row_offsets_devices=None)

source

Returns the named per-batch graph inputs, in the order fields gives.

Parameters:

• device_ref (DeviceRef) – The device that holds the token and row-offset tensors.
• fields (Sequence[str]) – The inputs to include, in graph order. Any of tokens, row_offsets, host_row_offsets, return_n_logits, data_parallel_splits and hidden_states.
• data_parallel_degree (int) – The number of replicas, which sizes data_parallel_splits.
• hidden_states (tuple[DType, int] | None) – The (dtype, hidden_size) of the EAGLE draft’s hidden-state input. Required when fields names hidden_states.
• row_offsets_devices (Sequence[DeviceRef] | None) – The devices that each take their own row_offsets, for a graph that does not broadcast them. row_offsets is then a tuple with one type per device.

Returns:

The input types keyed by field name.

Return type:

RaggedBatch