Scatter vectors [Solved]

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.