oneflow.vsplit

oneflow.vsplit()

Splits input, a tensor with two or more dimensions, into multiple tensors vertically according to indices_or_sections. Each split is a view of input. This is equivalent to calling oneflow.tensor_split(input, indices_or_sections, dim=0) (the split dimension is 0), except that if indices_or_sections is an integer it must evenly divide the split dimension or a runtime error will be thrown. The documentation is referenced from: https://pytorch.org/docs/1.10/generated/torch.vsplit.html.

Parameters
  • input (Tensor) – the input tensor.

  • indices_or_sections (int or a list) – If indices_or_sections is an integer n , input is split into n sections along dimension dim.If input is divisible by n along dimension dim, each section will be of equal size, input.size (dim) / n. If input is not divisible by n, the sizes of the first int(input.size(dim) % n). sections will have size int(input.size(dim) / n) + 1, and the rest will have size int(input.size(dim) / n). If indices_or_sections is a list or tuple of ints, then input is split along dimension dim at each of the indices in the list, tuple or tensor. For instance, indices_or_sections=[2, 3] and dim=0 would result in the tensors input[:2], input[2:3], and input[3:].If indices_or_sections is a tensor, it must be a zero-dimensional or one-dimensional long tensor on the CPU.

Returns

the output TensorTuple.

Return type

oneflow.TensorTuple

For example:

>>> import oneflow as flow

>>> input = flow.rand(4, 4, 5, 6)
>>> output = flow.vsplit(input, (1, 3))
>>> output[0].size()
oneflow.Size([1, 4, 5, 6])
>>> output[1].size()
oneflow.Size([2, 4, 5, 6])
>>> output[2].size()
oneflow.Size([1, 4, 5, 6])