oneflow.split

oneflow.split()

Splits the tensor into chunks.

If split_size_or_sections is an integer type, then x will be split into equally sized chunks (if possible). Last chunk will be smaller if the tensor size along the given dimension dim is not divisible by split_size.

If split_size_or_sections is a list, then x will be split into len(split_size_or_sections) chunks with sizes in dim according to split_size_or_sections.

Parameters
  • x – tensor to split.

  • split_size_or_sections – size of a single chunk or list of sizes for each chunk.

  • dim – dimension along which to split the tensor.

For example:

>>> import oneflow as flow
>>> a = flow.arange(10).view(5, 2)
>>> flow.split(a, 2)
(tensor([[0, 1],
        [2, 3]], dtype=oneflow.int64), tensor([[4, 5],
        [6, 7]], dtype=oneflow.int64), tensor([[8, 9]], dtype=oneflow.int64))
>>> flow.split(a, [1, 4])
(tensor([[0, 1]], dtype=oneflow.int64), tensor([[2, 3],
        [4, 5],
        [6, 7],
        [8, 9]], dtype=oneflow.int64))