split_data

split_data(data, num_slice, batch_axis=0, even_split=True)

Splits an NDArray into num_slice slices along batch_axis. Usually used for data parallelism where each slices is sent to one device (i.e. GPU).

Parameters:

data : NDArray

A batch of data.

num_slice : int

Number of desired slices.

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. If True, an error will be raised when num_slice does not evenly divide data.shape[batch_axis].

Returns:

list of NDArray

Return value is a list even if num_slice is 1.