tensor_split

→list[Tensor]

tensor_split(input: Tensor, indices_or_sections: _int | Sequence[_int], dim: _int = ...)

source edit

Split a tensor along dim, permitting unequal final-piece sizes.

More permissive than lucid.split: when indices_or_sections is an integer k and x.shape[dim] is not divisible by k, the first x.shape[dim] % k pieces receive one extra element each while the remaining pieces take the smaller floor size. This mirrors NumPy's np.array_split semantics, where split raises on non-divisible counts but tensor_split quietly returns ragged pieces.

Parameters

inputTensor

Input tensor.

indices_or_sectionsint | Sequence[int]

int: number of (near-)equal-sized splits. The first axis_len % k pieces receive one extra element.
Sequence[int]: cut indices along dim.

dimint

Axis to split along. Defaults to 0.

Returns

list[Tensor]

Sub-tensors whose concatenation along dim reproduces input.

Raises

lucid._C.engine.InvalidArgument

If input is 0-d, or indices_or_sections is an integer that is not positive.

Notes

For integer k with $n = \text{shape}(x)[\text{dim}]$ and $b = \lfloor n / k \rfloor$ , $r = n \bmod k$ , the piece sizes are

\underbrace{b + 1, \dots, b + 1}_{r\text{ terms}},\; \underbrace{b, \dots, b}_{k - r\text{ terms}}.

Contrast with lucid.split, whose integer argument is a piece size and whose size list must sum to n.

With a list of indices, piece $i$ is the slice x[indices[i-1]:indices[i]] along dim: a negative index counts from the end, an index past either end is clamped to it, and an index below the one before it gives an empty piece, with the next piece starting from that index again.

Examples

>>> import lucid
>>> x = lucid.arange(10)
>>> [s.shape for s in lucid.tensor_split(x, 3)]
[(4,), (3,), (3,)]
>>> [s.shape for s in lucid.tensor_split(x, [2, 5])]
[(2,), (3,), (5,)]
>>> [s.shape for s in lucid.tensor_split(x, [5, 2, 12])]
[(5,), (0,), (8,), (0,)]

>>> import lucid >>> x = lucid.arange(10) >>> [s.shape for s in lucid.tensor_split(x, 3)] [(4,), (3,), (3,)] >>> [s.shape for s in lucid.tensor_split(x, [2, 5])] [(2,), (3,), (5,)] >>> [s.shape for s in lucid.tensor_split(x, [5, 2, 12])] [(5,), (0,), (8,), (0,)]