chunk

→list of Tensor

chunk(input: Tensor, chunks: int, dim: DimLike = ...)

source

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

inputTensor

Source tensor.

chunksint

Desired number of pieces. Must be positive.

dimDimLike= ...

Dimension along which to split.

Returns

list of Tensor

Up to chunks views into input (no copy).

Notes

Each chunk has size $\lceil n / c \rceil$ 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,)]