Skip to content

Instantly share code, notes, and snippets.

@ptrblck
Created February 19, 2018 13:45
Show Gist options
  • Save ptrblck/a04b4ba8d641267746c2e85b3da6c899 to your computer and use it in GitHub Desktop.
Save ptrblck/a04b4ba8d641267746c2e85b3da6c899 to your computer and use it in GitHub Desktop.
import torch
import numpy as np
# numpy
a = np.random.rand(10, 20)
tmp0 = np.split(a, indices_or_sections=5, axis=0) # split into 5 sections
for t in tmp0:
print(t.shape)
np.split(a, indices_or_sections=7, axis=0) # error, since no equal division
tmp1 = np.split(a, [5, 7], 0) # use indices ([:5], [5:7], [7:])
for t in tmp1:
print(t.shape)
# PyTorch
x = torch.randn(10, 20)
tmp2 = torch.split(x, split_size_or_sections=4, dim=0) # use size 4
for t in tmp2:
print(t.shape) # last split might be smaller
tmp3 = torch.split(x, split_size_or_sections=[5, 2, 3], dim=0)
for t in tmp3:
print(t.shape)
torch.split(x, split_size_or_sections=[5, 4], dim=0) # error, since 5+4 != dim(0)
# Should it return Tensors of size [5, 20], [4, 20] and [1, 20]?
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment