megengine.functional.split#

split(inp, nsplits_or_sections, axis=0)[source]#

Splits the input tensor into several smaller tensors. When nsplits_or_sections is int, the last tensor may be smaller than others.

Parameters:

  • inp – input tensor.

  • nsplits_or_sections – number of sub tensors or sections information list.

  • axis – which axis will be splited.

Returns:

output tensor list.

Examples

>>> import os
>>> import numpy as np
>>> x = Tensor(np.random.random((10, 20)), dtype=np.float32)
>>> y = F.split(x, 3)
>>> z = F.split(x, [6, 17], axis=1)
>>> print([i.numpy().shape for i in y])
[(4, 20), (3, 20), (3, 20)]
>>> print([i.numpy().shape for i in z])
[(10, 6), (10, 11), (10, 3)]