Repeating tensor rows given lengths tensor

Hello!

Let’s say I have a length’s tensor (lengths_tensor) : len of shape [32 X count] for example :

tensor([1, 1, 1, 1, 3, 1, 1, 2, 1, 2, 3, 1, 2, 2, 2, 3, 4, 1, 3, 1, 2, 2, 2, 2, 2, 2, 1, 1, 2, 2, 2, 1], device=‘cuda:0’)

and another features tensor of shape [32 X Dimension]

what I want to do is for each entry (i) in the features tensor to repeat it lengths_tensor[i] times, so given the example above the first element in the feature tensor should not be repeated, and for the element in place 4 shall be repeated 3 times, so the new feature tensor should be of shape [60 x Dimension]
Thanks!

I’m unsure how you’ve calculated 60 elements, as the sum if your lengths_tensor is 57.
Assuming 57 is the expected shape, then repeat_interleave might work:

a = torch.tensor([1, 1, 1, 1, 3, 1, 1, 2, 1, 2, 3, 1, 2, 2, 2, 3, 4, 1, 3, 1, 2, 2, 2, 2, 2, 2, 1, 1, 2, 2, 2, 1])

x = torch.arange(32)
y = x.repeat_interleave(a)
print(y)
# tensor([ 0,  1,  2,  3,  4,  4,  4,  5,  6,  7,  7,  8,  9,  9, 10, 10, 10, 11,
#         12, 12, 13, 13, 14, 14, 15, 15, 15, 16, 16, 16, 16, 17, 18, 18, 18, 19,
#         20, 20, 21, 21, 22, 22, 23, 23, 24, 24, 25, 25, 26, 27, 28, 28, 29, 29,
#         30, 30, 31])