oneflow.tensor_split¶
-
oneflow.
tensor_split
()¶ Splits a tensor into multiple sub-tensors, all of which are views of input, along dimension dim according to the indices or number of sections specified by indices_or_sections . The documentation is referenced from: https://pytorch.org/docs/1.10/generated/torch.tensor_split.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.
dim (int) – dimension along which to split the tensor.
- Returns
the output TensorTuple.
- Return type
oneflow.TensorTuple
For example:
>>> import oneflow as flow >>> input = flow.rand(3,4,5) >>> output = flow.tensor_split(input,(2,3),2) >>> output[0].size() oneflow.Size([3, 4, 2]) >>> output[1].size() oneflow.Size([3, 4, 1]) >>> output[2].size() oneflow.Size([3, 4, 2])