Implementing kernel
C++ engine symbols that back this Python API.Split a tensor into at most chunks pieces along dim.
Each piece has size along dim except the
last, which holds what remains — so fewer than chunks pieces come
back when that size exhausts the dimension early (n=6, c=4 gives
[2, 2, 2]). split takes the piece size instead.
Parameters
inputTensorSource tensor.
chunksintDesired number of pieces. Must be positive.
dimint= 0Dimension along which to split.
Returns
list of TensorThe pieces of input in order — views of its buffer on the CPU,
copies on metal.
Raises
lucid._C.engine.InvalidArgumentchunks is not positive.Examples
>>> import lucid
>>> x = lucid.arange(7)
>>> [t.shape for t in lucid.chunk(x, 3, dim=0)]
[(3,), (3,), (1,)]
>>> [t.shape for t in lucid.chunk(x, 4, dim=0)]
[(2,), (2,), (2,), (1,)]