mindspore.mint.chunk

View Source On Gitee
mindspore.mint.chunk(input, chunks, dim=0)[source]

Cut the input Tensor into chunks sub-tensors along the specified axis.

Note

This function may return less than the specified number of chunks!

Warning

This is an experimental API that is subject to change or deletion.

Parameters

  • input (Tensor) – A Tensor to be cut.

  • chunks (int) – Number of sub-tensors to cut.

  • dim (int, optional) – Specify the dimensions that you want to split. Default: 0 .

Returns

A tuple of sub-tensors.

Raises

  • TypeError – If argument input is not Tensor.

  • TypeError – The sum of chunks is not int.

  • TypeError – If argument dim is not int.

  • ValueError – If argument dim is out of range of \([-input.ndim, input.ndim)\) .

  • ValueError – If argument chunks is not positive number.

Supported Platforms:

Ascend

Examples

>>> import numpy as np
>>> import mindspore
>>> from mindspore import Tensor
>>> input_x = np.arange(9).astype("float32")
>>> output = mindspore.mint.chunk(Tensor(input_x), 3)
>>> print(output)
(Tensor(shape=[3], dtype=Float32, value= [ 0.00000000e+00,  1.00000000e+00,  2.00000000e+00]),
 Tensor(shape=[3], dtype=Float32, value= [ 3.00000000e+00,  4.00000000e+00,  5.00000000e+00]),
 Tensor(shape=[3], dtype=Float32, value= [ 6.00000000e+00,  7.00000000e+00,  8.00000000e+00]))