split(tensor: Tensor, split_size_or_sections: _int | list[_int], dim: _int = 0)Implementing kernel
C++ engine symbols that back this Python API.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
tensorTensorSource tensor.
split_size_or_sectionsint or list of intPer-piece size, or a list of sizes.
dimint= 0Dimension to split along.
Returns
list of TensorThe pieces of tensor — views of its buffer on the CPU, copies on
metal.
Raises
lucid._C.engine.InvalidArgumentA negative size, a zero integer size on a non-empty axis, or a size
list that does not sum to
input.size(dim) — the last piece never
takes up a remainder or runs short. It is both a ValueError and
a RuntimeError.TypeErrorA size that is not an integer.
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,)]