split_and_load

split_and_load(data, ctx_list, batch_axis=0, even_split=True)

Splits an NDArray into len(ctx_list) slices along batch_axis and loads each slice to one context in ctx_list.

Parameters:

data : NDArray or ndarray

A batch of data.

ctx_list : list of Context

A list of Contexts.

batch_axis : int, default 0

The axis along which to slice.

even_split : bool, default True

Whether to force all slices to have the same number of elements.

Returns:

list of NDArrays or ndarrays

Each corresponds to a context in ctx_list.