split

→list of Tensor

split(x: Tensor, split_size_or_sections: _int | list[_int], dim: _int = 0)

source

Split a tensor into pieces of explicit sizes along dim.

If split_size_or_sections is an integer, the tensor is split into equal pieces of that size (except possibly the last, which may be smaller). If it is a list, each entry specifies the size of one output piece and the list must sum to input.size(dim).

Parameters

xTensor

Source tensor.

split_size_or_sectionsint or list of int

Per-piece size, or a list of sizes.

dimint= 0

Dimension to split along.

Returns

list of Tensor

Views into x — no data is copied.

Notes

For "as many pieces as fit" semantics with no remainder constraint, prefer chunk.

Examples

>>> import lucid
>>> x = lucid.arange(7)
>>> [t.shape for t in lucid.split(x, [2, 3, 2])]
[(2,), (3,), (2,)]