Hi all,
I was wondering if there was a smart pytorch-ic way to extract the rows of a 3D tensor based on some 3D ‘span’ index tensor.
In detail, let’s say you have a (32, 20, 768) tensor x, where 32 is the batchsize, 20 is the sequence length (sentence length for instance) and 768 the encoding dimension. You also have an index tensor i of shape (32, 3, 2), where again the first dimension is the batchsize (corresponding to x), the second dimension is the number of spans you want to extract and the last one is equal to 2 since a span is defined by two integers.
For example let’s say that i[0] is:

[ [1,2], [5,7], [15,19] ]

what I would like to extract from x[0] is:

x[0][1:2]
x[0][5:7]
x[0][15:19]

and build a new tensor y with y[0] of the form:

[ x[0][1:2].mean(0),
x[0][5:7].mean(0),
x[0][15:19].mean(0)) ]

and in general

y[s] = [ x[s][ i[s][u][0] : i[s][u][1] ].mean(0) ]

with s in range(32) and u in range(3) for each s.
I guess the solution should be somewhat similar to what discussed here 2, but I can’t figure out how to achieve the desired result avoiding for loops.