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)

source

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

The result is an OrderedDict, not a dict: tree.flatten sorts a plain dict’s keys, and the graph declares its inputs in insertion order.

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.

Returns:

The input types keyed by field name.

Return type:

OrderedDict[str, TensorType]