Multiple sums based on index

Fst_Msn: <p>Given a 3D tensor A of size</p>
<blockquote>
<p>(n, i, j)</p>
</blockquote>
<p>And an index tensor of size n with a pattern like:</p>
<blockquote>
<p>[1, 2, 2, 3, 3, 3, 4, 4, 4, 4 …]</p>
</blockquote>
<p>I want to sum over the n dimensions of A according to the index tensor.</p>
<p>For example, if A is (6, 2, 2)</p>
<pre><code class="lang-auto">tensor([[[0.4, 0.4],
         [0.6, 0.3]],

        [[0.9, 0.3],
         [0.5, 0.6]],

        [[0.4, 0.1],
         [0.4, 0.8]],

        [[0.6, 0.3],
         [0.7, 0.1]],

        [[0.0, 0.9],
         [0.4, 0.7]],

        [[0.0, 0.4],
         [0.5, 0.7]]])
</code></pre>
<p>I want to sum A[1] + A[2] and A[3] + A[4] + A[5].<br>
The result has size (3,2,2) and is equal to:</p>
<pre><code class="lang-auto">tensor([[[0.4, 0.4],
        [0.6, 0.3]],

       [[1.3, 0.4],
        [0.9, 1.4]],

       [[0.6, 1.6],
        [1.6, 1.5]]])
</code></pre>
<p>Which is the best way to do this, for large n ?</p>

KFrank: <p>Hi Fst!</p>
<aside class="quote no-group" data-username="Fst_Msn" data-post="1" data-topic="161796" data-full="true">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/fst_msn/48/31467_2.png" class="avatar"> Fst_Msn:</div>
<blockquote>
<p>I want to sum A[1] + A[2] and A[3] + A[4] + A[5].<br>
The result has size (3,2,2) and is equal to:</p>
<pre><code class="lang-auto">tensor([[[0.4, 0.4],
        [0.6, 0.3]],

       [[1.3, 0.4],
        [0.9, 1.4]],

       [[0.6, 1.6],
        [1.6, 1.5]]])
</code></pre>
<p>Which is the best way to do this, for large n ?</p>
</blockquote>
</aside>
<p>For efficiency for large <code>n</code>, the best way will be to compute the<br>
cumulative sum of <code>A</code> (along the <code>0</code> dimension) and then obtain<br>
the desired partial sums as differences of the appropriate cumulative<br>
sums.</p>
<p>The cumulative sum requires (about) <code>n</code> additions, but you have<br>
to perform (about) <code>n</code> additions in any event, so packaging the<br>
computation as a cumulative sum probably lets pytorch make the<br>
most efficient use of pipelining.</p>
<p>The rest is just monkeying around to index into the cumulative sum.</p>
<p>Consider:</p>
<pre><code class="lang-plaintext">&gt;&gt;&gt; import torch
&gt;&gt;&gt; print (torch.__version__)
1.12.0
&gt;&gt;&gt;
&gt;&gt;&gt; # example tensor
&gt;&gt;&gt; A = torch.tensor ([[[0.4, 0.4],
...                     [0.6, 0.3]],
...                    [[0.9, 0.3],
...                     [0.5, 0.6]],
...                    [[0.4, 0.1],
...                     [0.4, 0.8]],
...                    [[0.6, 0.3],
...                     [0.7, 0.1]],
...                    [[0.0, 0.9],
...                     [0.4, 0.7]],
...                    [[0.0, 0.4],
...                     [0.5, 0.7]]])
&gt;&gt;&gt;
&gt;&gt;&gt; n = A.size (0)
&gt;&gt;&gt;
&gt;&gt;&gt; # assume that n is a "triangle number"
&gt;&gt;&gt; k = int ((2 * n)**0.5)
&gt;&gt;&gt;
&gt;&gt;&gt; seq = torch.arange (1, k + 1)
&gt;&gt;&gt; ind_up = ((seq * (seq + 1)) / 2).long() - 1
&gt;&gt;&gt; ind_lo = ((seq * (seq - 1)) / 2).long() - 1   # incorrect negative index will be masked out
&gt;&gt;&gt; msk = (ind_lo &gt;= 0).float()  # mask to get correct zero-th slice of result
&gt;&gt;&gt;
&gt;&gt;&gt; Acs = A.cumsum (0)   # cumulative sum of slices
&gt;&gt;&gt; res = Acs[ind_up] - msk[:, None, None] * Acs[ind_lo]   # clever differences of cumulative sums
&gt;&gt;&gt;
&gt;&gt;&gt; res   # final result
tensor([[[0.4000, 0.4000],
         [0.6000, 0.3000]],

        [[1.3000, 0.4000],
         [0.9000, 1.4000]],

        [[0.6000, 1.6000],
         [1.6000, 1.5000]]])
</code></pre>
<p>Best.</p>
<p>K. Frank</p>