[FIXED] Creating a new tensor according to a list of lengths

Issue

I have a tensor t with dim b x 3 and a list of lengths len = [l_0, l_1, ..., l_n]. All entries in len sum to to b. I want to create a new tensor with dim n x 3, which stores the average of the entries in t. E.g. The first l_0 entries in t are averaged and build the first element in the new tensor. The following l_1 entries are averaged and build the second element, ...

Thanks for your help.

Solution

You can do so using a combination a cumulative list of indices as helper and a list comprehension to construct the new array:

>>> b, lens = 10, [2, 3, 1, 3, 1]

>>> t =  torch.rand(b, 3)
tensor([[0.3567, 0.3998, 0.9396],
        [0.4061, 0.6465, 0.6955],
        [0.3500, 0.4135, 0.5288],
        [0.0726, 0.9575, 0.3785],
        [0.6216, 0.2975, 0.3293],
        [0.3878, 0.0735, 0.8181],
        [0.1694, 0.5446, 0.1179],
        [0.7793, 0.6613, 0.1748],
        [0.0964, 0.9825, 0.1651],
        [0.1421, 0.0994, 0.8086]])

Build the cumulative list of indices:

>>> c = torch.cumsum(torch.tensor([0] + lens), 0)
tensor([ 0,  2,  5,  6,  9, 10])

Loop over c by twos, with an overlapping window. For example zip(c[:-1], c[1:]) works well. Each selection from i to j gets averaged on dim=0.

>>> [t[i:j].sum(0) for i, j in zip(c[:-1], c[1:])]
[tensor([0.7628, 1.0463, 1.6351]),
 tensor([1.0442, 1.6685, 1.2367]),
 tensor([0.3878, 0.0735, 0.8181]),
 tensor([1.0451, 2.1885, 0.4578]),
 tensor([0.1421, 0.0994, 0.8086])]

Then you can stack the list:

>>> torch.stack([t[i:j].sum(0) for i, j in zip(c[:-1], c[1:])])
tensor([[0.7628, 1.0463, 1.6351],
        [1.0442, 1.6685, 1.2367],
        [0.3878, 0.0735, 0.8181],
        [1.0451, 2.1885, 0.4578],
        [0.1421, 0.0994, 0.8086]])

Answered By - Ivan