Reshape and pad a tensor given a list of lengths
I have given a 2d tensor in of shape a x b like the following (where a = 9 and each of A1, A2, ..., C2 represents a b-dimensional vector): 



Furthermore, I have an array of lengths, where sum(lengths) = a and each entry is a positive integer:



Then I would like to obtain a 3d output tensor out, where the first lengths[0] entries of in form the first row, the next lengths[1] entries of in form the second row, and so on. That is, the output tensor should have the shape len(lengths) x max(lengths) x b, and be padded with zeros (each 0 in the below picture represents a b-dimensional zero vector):



As this is part of a neural network that is trained using backpropagation, all operations used must be differentiable. How can this be achieved (ideally, with good performance) using PyTorch?

Here is my implementation using torch.nn.utils.rnn.pad_sequence():

in_tensor = torch.rand((9, 3))
print(in_tensor)
print(36*'=')
lengths = torch.tensor([3, 4, 2])
cum_len = 0
y = []
for idx, val in enumerate(lengths):
    y.append(in_tensor[cum_len : cum_len+val])
    cum_len += val
print(torch.nn.utils.rnn.pad_sequence(y, batch_first=True)))


output:

# in_tensor of shape (9 x 3)
tensor([[0.9169, 0.3549, 0.6211],
        [0.4832, 0.5475, 0.8862],
        [0.8708, 0.5462, 0.9374],
        [0.4605, 0.1167, 0.5842],
        [0.1670, 0.2862, 0.0378],
        [0.2438, 0.5742, 0.4907],
        [0.1045, 0.5294, 0.5262],
        [0.0805, 0.2065, 0.2080],
        [0.6417, 0.4479, 0.0688]])
====================================
# out tensor of shape (len(lengths) x max(lengths) x b), in this case b is 3
tensor([[[0.9169, 0.3549, 0.6211],
         [0.4832, 0.5475, 0.8862],
         [0.8708, 0.5462, 0.9374],
         [0.0000, 0.0000, 0.0000]],

        [[0.4605, 0.1167, 0.5842],
         [0.1670, 0.2862, 0.0378],
         [0.2438, 0.5742, 0.4907],
         [0.1045, 0.5294, 0.5262]],

        [[0.0805, 0.2065, 0.2080],
         [0.6417, 0.4479, 0.0688],
         [0.0000, 0.0000, 0.0000],
         [0.0000, 0.0000, 0.0000]]])

