split_with_sizes

paddle.Tensor. split_with_sizes ( self: Tensor, split_sizes: list[int], dim: int = 0 ) → list[paddle.Tensor]

Splits the input tensor into multiple sub tensors according to given split sizes.

Parameters

• self (Tensor) – The input tensor to be split.

• split_sizes (list[int]) – A list of non negative integers specifying the sizes of each split along dimension dim. The sum of all elements in this list must equal the size of self along dim.

• dim (int, optional) – The dimension along which to split the tensor. Defaults to 0.

Returns

A list of sub tensors resulting from splitting self along the specified dimension.

Return type

list[Tensor]

Examples

>>> import paddle
>>> x = paddle.to_tensor([[1, 2, 3, 4], [5, 6, 7, 8], [9, 10, 11, 12]])
>>> # Split into two parts along the first dimension, of sizes 1 and 2
>>> splits = paddle.Tensor.split_with_sizes(x, [1, 2], dim=0)
>>> print(splits)