Python function

split_batch

`split_batch()`

max.nn.split_batch(devices, input, input_row_offsets, data_parallel_splits)

source

Split a ragged input batch into data parallel batches.

devices = [device_1, device_2]
input = [seq_1, seq_2, seq_3, seq_4]
input_row_offsets = [0, offset_1, offset_2, offset_3, offset_4]
data_parallel_splits = [0, 2, 4]

# Outputs
split_input = [seq_1, seq_2], [seq_3, seq_4]
split_offsets = [0, offset_1, offset_2], [0, new_offset_3, new_offset_4]

This method places the outputs on the devices specified in devices.

See split_batch_replicated() for a version of this method that takes replicated inputs and input_row_offsets for each device.

Parameters:

input (TensorValue) – Input tensor of shape [total_seq_len, …].
input_row_offsets (TensorValue) – Row offsets tensor indicating batch boundaries.
data_parallel_splits (TensorValue) – Buffer containing batch splits for each device that must be located on CPU. The size of data_parallel_splits must be equal to the number of devices + 1.
devices (list[DeviceRef])

Returns:

Tuple of (split_input, split_offsets) where split_input and split_offsets are lists of tensors, one per device

Return type:

tuple[list[TensorValue], list[TensorValue]]

split_batch()​

`split_batch()`