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

seyeeet: <p>I have a tensor <code>A=torch.rand(1,512,100)</code> where <code>100</code> is the duration/time in my tensor.<br>
I have another variable <code>T</code> that shows the duration for me, e.g., <code>T=torch.tensor([[20,30,40,10]])</code>.<br>
I would like to compute the average of the elements in <code>A</code> based on the durations in <code>T</code> so the output will be in shape of <code>A.shape[0],A,shape[1],T.shape[2]</code>.<br>
in other words, I want to do:</p>
<pre><code class="lang-auto">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)
</code></pre>
<p>is there a way to do it (without for loops) so I can do it automatically for when <code>T</code> has different length in dimension 1? i.e. one time <code>T</code> can be  <code>T=torch.tensor([[20,30,40,10]])</code> and another time <code>T</code> can be  <code>T=torch.tensor([[50,50]])</code>, the sum of <code>T</code> elements are always equal to <code>A.shape[2]</code></p>

zetyquickly: <p>HI <a class="mention" href="/u/seyeeet">@seyeeet</a></p>
<p>First, if <code>T=torch.tensor([20,30,40,10])</code> then <code>A[:,:,T[2]:T[3]]</code> won’t give you anything, because <code>10</code> is less then <code>40</code>. Thus I propose to use <code>cumsum</code> first to get <code>torch.tensor([20,50,90,100])</code></p>
<p>Second, you can use the resulting tensor as the index tensor for the <code>A</code>. But you still will need a loop to get <code>A[T[i]]</code> because the slices will be of different 3d sizes</p>
<p>Btw, I don’t exclude it that there’s a possibility to create 4d tensor that when apply <code>mean(axis=3)</code> would give what you’re looking for</p>

seyeeet: <p>Hi <a class="mention" href="/u/zetyquickly">@zetyquickly</a><br>
Thanks for pointing that out, yes you are correct, it should be mean on <code>cumsum</code>, I corrected the post regarding that.</p>
<p>Yes, doing the loop over the index will gives the results.</p>
<p>I wonder if there is a non-loop version that I am not aware of it because loops can be slow for large length.</p>
<p>I am not sure if I got the <code>mean(axis=3)</code> suggestion. can you please elaborate on that?</p>
<p>Thanks for your help! <img src="https://discuss.pytorch.org/images/emoji/apple/slight_smile.png?v=10" title=":slight_smile:" class="emoji" alt=":slight_smile:"></p>

zetyquickly: <p>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 <code>.mean(axis=3)</code> on it. As a result it would be 3d tensor.</p>
<p>Try to think in that way, which 4d tensor you can construct first</p>

seyeeet: <p>hi <a class="mention" href="/u/ptrblck">@ptrblck</a><br>
do you by any chance have a more pytorchic suggestion that can do this task in a faster and more general way?</p>