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))