Compute the sum of first k[i] elements of a tensor

username: youzunzhi
cooked: <p>Suppose I have a 2-d tensor:</p>
<pre><code class="lang-auto">t
tensor([[1.0000, 3.0000, 4.0000, 2.0000, 1.0000],
        [2.0000, 3.0000, 3.5000, 4.0000, 3.5000]])
</code></pre>
<p>and a LongTensor:</p>
<pre><code class="lang-auto">c
tensor([2, 3])
</code></pre>
<p>I want to get the first <code>c[i]</code> elements’ sum of the i-th row of <code>t</code>, i.e. [4.0, 8.5] in this case. How to calculate this without iteration?</p>

username: Dexter
cooked: <p>You have a 2-row tensor and your index tensor is <code>[2, 3]</code>, so I think what you are asking is a <code>t[:, 2].sum()</code> and <code>t[:, 3].sum()</code>, correct me if I am wrong, if that is so then how are you getting <code>[4.0, 8.5]</code>?</p>

username: youzunzhi
cooked: <p>Sorry, but I am trying to get  <code>t[0, :2].sum()</code>  and  <code>t[1, :3].sum()</code> , which is 4.0 and 8.5 respectively. Thank you for your reply.</p>

username: youzunzhi
cooked: <p>Now I have a solution:</p>
<pre><code class="lang-auto">t.cumsum(1)[torch.arange(len(c)), c-1]
tensor([4.0000, 8.5000])
</code></pre>
<p>But it does not work if <code>c[i]</code> is 0 for any i. And it does a lot of redundant calculation in the <code>cumsum()</code>. So maybe there is a better solution?</p>

username: Dexter
cooked: <p>Have you checked <code>torch.narrow</code>, you can get a slice of your tensor then apply <code>sum(axis=1)</code> but the elements in tensor c should have the same value else it won’t return a result. Other than that looks like you really need to use loops.</p>

username: youzunzhi
cooked: <p>Thanks, but it seems  <code>torch.narrow</code>  cannot slice different length of each row.</p>

username: youzunzhi
cooked: <p>Inspired by a friend Zijing Ou, here is an easy solution by constructing a mask:</p>
<pre><code class="lang-auto">mask = torch.linspace(0, t.shape[1]-1, t.shape[1]).view(1, t.shape[1]).repeat(t.shape[0], 1) &lt; c.view(t.shape[0], 1)
(t * mask).sum(1)
</code></pre>
<p>would get the desired results.</p>

username: Dexter
cooked: <p>Thanks for sharing. I’ll note this one.</p>

username: KFrank
cooked: <p>Hello Youzunzhi!</p>
<aside class="quote no-group" data-username="youzunzhi" data-post="1" data-topic="96905" data-full="true">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/letter_avatar_proxy/v4/letter/y/a698b9/48.png" class="avatar"> youzunzhi:</div>
<blockquote>
<p>I want to get the first <code>c[i]</code> elements’ sum of the i-th row of <code>t</code></p>
</blockquote>
</aside>
<p>I don’t know of any way to avoid a loop without summing over the<br>
full tensor, such as by using <code>.cumsum()</code> (or Zijing’s masked-sum<br>
suggestion).</p>
<p>After forming the sums, you have to index into them to get the<br>
desired partial sum.  I don’t know of a “clean” way of doing this;<br>
the best I could come up with is to use <code>.take()</code>.</p>
<p>To accommodate the case in which an element of <code>c</code> is <code>0</code> (so that<br>
you sum over zero elements of <code>t</code>, getting <code>0.0</code>), we “initialize” the<br>
partial sums with <code>0.0</code> by prepending a zero slice to <code>t</code>.</p>
<p>Here is my approach:</p>
<pre data-code-wrap="python"><code class="lang-python">import torch
torch.__version__

t = torch.FloatTensor([[1.0000, 3.0000, 4.0000, 2.0000, 1.0000],
                       [2.0000, 3.0000, 3.5000, 4.0000, 3.5000]])
c = torch.LongTensor([2, 3])

t0 = torch.cat ((torch.zeros (t.shape[0], 1), t), 1)      # initialize partial sums with 0
tsum = t0.cumsum (dim = 1)                                # calculate partial sums
cp = c + t0.shape[1] * torch.arange (c.shape[0]).long()   # 1-d indices so we can use take()
sc = tsum.take (cp)                                       # get specified partial sums

print ('sc =\n', sc)                                      # print result
</code></pre>
<p>And here is the output:</p>
<pre><code class="lang-plaintext">&gt;&gt;&gt; import torch
&gt;&gt;&gt; torch.__version__
'0.3.0b0+591e73e'
&gt;&gt;&gt;
&gt;&gt;&gt; t = torch.FloatTensor([[1.0000, 3.0000, 4.0000, 2.0000, 1.0000],
...                        [2.0000, 3.0000, 3.5000, 4.0000, 3.5000]])
&gt;&gt;&gt; c = torch.LongTensor([2, 3])
&gt;&gt;&gt;
&gt;&gt;&gt; t0 = torch.cat ((torch.zeros (t.shape[0], 1), t), 1)      # initialize partial sums with 0
&gt;&gt;&gt; tsum = t0.cumsum (dim = 1)                                # calculate partial sums
&gt;&gt;&gt; cp = c + t0.shape[1] * torch.arange (c.shape[0]).long()   # 1-d indices so we can use take()
&gt;&gt;&gt; sc = tsum.take (cp)                                       # get specified partial sums
&gt;&gt;&gt;
&gt;&gt;&gt; print ('sc =\n', sc)                                      # print result
sc =

 4.0000
 8.5000
[torch.FloatTensor of size 2]
</code></pre>
<p>Best.</p>
<p>K. Frank</p>