Implementing kernel
C++ engine symbols that back this Python API.Split a tensor into chunks roughly equal pieces along dim.
Unlike split, which takes an explicit size or list of sizes,
chunk divides the input dimension into chunks parts whose
sizes differ by at most one element. If dim is not divisible by
chunks the last chunk is smaller.
Parameters
inputTensorSource tensor.
chunksintDesired number of pieces. Must be positive.
dimDimLike= ...Dimension along which to split.
Returns
list of TensorUp to chunks pieces of input — views of its buffer on the
CPU, copies on metal.
Notes
Each chunk has size except possibly the
last; for n=7, c=3 the chunks are [3, 3, 1].
Examples
>>> import lucid
>>> x = lucid.arange(7)
>>> [t.shape for t in lucid.chunk(x, 3, dim=0)]
[(3,), (3,), (1,)]