EDIT: Forgot about advanced indexing, solution below:
num_vectors = 5
max_length = 3
num_out = 3
t = torch.arange(num_vectors).view(-1, 1).float()
zeros = torch.zeros(num_out * max_length, 1)
zeros[[0, 1, 2, 3, 6, 7]] = t[[2, 1, 0, 4, 3, 2]]
print(zeros.view(3, 3, 1))
tensor([[[2.],
[1.],
[0.]],
[[4.],
[0.],
[0.]],
[[3.],
[2.],
[0.]]])
OLD POST:
Hey,
I want to scatter vectors according to a list of indices of variable lengths while padding with zeros like this:
vectors = torch.tensor(
[
[5],
[6],
[7]
]
)
# vectors.size(): (num_vectors, vector_size) = (3, 1)
indices = [[1, 2], [0], [0, 1, 1, 1]]
scattered = scatter(vectors, indices, padding=0)
# scattered:
# torch.tensor(
# [
# [[6], [7], [0], [0]],
# [[5], [0], [0], [0]],
# [[5], [6], [6], [6]],
# ]
# )
I am currently using tensor._scatter
like this:
def padded_scatter(src, indices, padding=0):
num_items = [len(idx) for idx in indices]
n_max = max(num_items)
out = torch.zeros(
n_max * len(num_items),
*src.size()[1:],
dtype=src.dtype,
device=src.device
)
if padding:
out[:] = padding
scatter_index = torch.zeros(src.size(), dtype=torch.int64, device=src.device)
for g_i, group_indices in enumerate(indices):
for i, idx in enumerate(group_indices):
scatter_index[idx] = g_i * n_max + i
out.scatter_(0, scatter_index, src)
return out.view(len(num_items), n_max, *src.size()[1:]), num_items
I want to scatter entire vectors, is there a better way to do that? Seems inefficient to scatter each element by itself.
I could always copy each vector in a loop but I’d like to avoid that as it seems pretty slow on CUDA to do lots of smaller actions in a loop.