chunk

paddle. chunk ( x: Tensor, chunks: int, axis: int | Tensor = 0, name: str | None = None ) → list[Tensor] [source]

Split the input tensor into multiple sub-Tensors.

Here are some examples to explain it.

    1. Given a 3-D tensor x with a shape [3, 3, 3], if we split the first dimension into three equal parts, it will output a list containing three 3-D tensors with a shape of [1, 3, 3].

    1. Given a 3-D tensor x with a shape [3, 3, 3], if we split the second dimension into three equal parts, it will output a list containing three 3-D tensors with a shape of [3, 1, 3].

    1. Given a 3-D tensor x with a shape [3, 3, 3], if we split the third dimension into three equal parts, it will output a list containing three 3-D tensors with a shape of [3, 3, 1].

The following figure illustrates the first example.

legend of reshape API

Note

Alias Support: The parameter name input can be used as an alias for x, and the parameter name dim can be used as an alias for axis. For example, chunk(input=tensor_x, dim=1) is equivalent to chunk(x=tensor_x, axis=1).

Parameters

  • x (Tensor) – A N-D Tensor. The data type is bool, float16, float32, float64, int32 or int64. alias: input.

  • chunks (int) – The number of tensor to be split along the certain axis.

  • axis (int|Tensor, optional) – The axis along which to split, it can be a integer or a 0-D Tensor with shape [] and data type int32 or int64. If :math::axis < 0, the axis to split along is \(rank(x) + axis\). Default is 0. alias: dim.

  • name (str|None, optional) – The default value is None. Normally there is no need for user to set this property. For more information, please refer to api_guide_Name .

Returns

list(Tensor), The list of segmented Tensors.

Examples

>>> import paddle

>>> x = paddle.rand([3, 9, 5])

>>> out0, out1, out2 = paddle.chunk(x, chunks=3, axis=1)
>>> # out0.shape [3, 3, 5]
>>> # out1.shape [3, 3, 5]
>>> # out2.shape [3, 3, 5]


>>> # axis is negative, the real axis is (rank(x) + axis) which real
>>> # value is 1.
>>> out0, out1, out2 = paddle.chunk(x, chunks=3, axis=-2)
>>> # out0.shape [3, 3, 5]
>>> # out1.shape [3, 3, 5]
>>> # out2.shape [3, 3, 5]