How to compute mean of a tensor based on a selective range span?

1

I have a tensor A=torch.rand(1,512,100) where 100 is the duration/time in my tensor.
I have another variable T that shows the duration for me, e.g., T=torch.tensor([[20,30,40,10]]).
I would like to compute the average of the elements in A based on the durations in T so the output will be in shape of A.shape[0],A,shape[1],T.shape[2].
in other words, I want to do:

A=torch.rand(1,512,100)
T=torch.tensor([[20,30,40,10]])
T_cumsum =torch.cumsum(T,dim=1)
results = torch.zeros(A.shape[0],A.shape[1],T.shape[1])
results[:,:,0] = A[:,:,0:T_cumsum [0,0]].mean(2)
results[:,:,1] = A[:,:,T_cumsum [0,0]:T_cumsum [0,3]].mean(2)
results[:,:,2] = A[:,:,T_cumsum [0,1]:T_cumsum [0,2]].mean(2)
results[:,:,3] = A[:,:,T_cumsum [0,2]:].mean(2)

is there a way to do it (without for loops) so I can do it automatically for when T has different length in dimension 1? i.e. one time T can be T=torch.tensor([[20,30,40,10]]) and another time T can be T=torch.tensor([[50,50]]), the sum of T elements are always equal to A.shape[2]

2

HI @seyeeet

First, if T=torch.tensor([20,30,40,10]) then A[:,:,T[2]:T[3]] won’t give you anything, because 10 is less then 40. Thus I propose to use cumsum first to get torch.tensor([20,50,90,100])

Second, you can use the resulting tensor as the index tensor for the A. But you still will need a loop to get A[T[i]] because the slices will be of different 3d sizes

Btw, I don’t exclude it that there’s a possibility to create 4d tensor that when apply mean(axis=3) would give what you’re looking for

1 Like
3

Hi @zetyquickly
Thanks for pointing that out, yes you are correct, it should be mean on cumsum, I corrected the post regarding that.

Yes, doing the loop over the index will gives the results.

I wonder if there is a non-loop version that I am not aware of it because loops can be slow for large length.

I am not sure if I got the mean(axis=3) suggestion. can you please elaborate on that?

Thanks for your help! :slight_smile:

4

I mean that, if you are expecting to get the 3d output tensor of means via matrix operations, then it is highly likely that you should construct 4d tensor first and then run .mean(axis=3) on it. As a result it would be 3d tensor.

Try to think in that way, which 4d tensor you can construct first

5

hi @ptrblck
do you by any chance have a more pytorchic suggestion that can do this task in a faster and more general way?