Pytorch equivalent to `tf.unsorted_segment_sum`

handesy: <p>Hi,</p>
<p>I was wondering what the equivalent function to <code>tf.unsorted_segment_sum</code> (<a href="https://www.tensorflow.org/api_docs/python/tf/unsorted_segment_sum" rel="nofollow noopener">https://www.tensorflow.org/api_docs/python/tf/unsorted_segment_sum</a>) is. I can come up with two alternatives:</p>
<ol>
<li>Use (sparse) matrix multiplication. Suppose <code>Y</code> has shape <code>(M,D)</code> and <code>X</code> has shape <code>(N, D)</code>, we can do</li>
</ol>
<pre><code class="lang-auto">Y = I @ X
</code></pre>
<p>where <code>I</code> is a <code>M</code> by <code>N</code> sparse binary matrix. The issue with this implementation is the overhead of constructing <code>I</code>.</p>
<ol start="2">
<li>Use <code>scatter_add_</code>, like</li>
</ol>
<pre><code class="lang-auto">Y = torch.zeros(M, D).scatter_add_(dim=0, index=I, other=X)
</code></pre>
<p>where <code>I</code> is an index matrix of shape <code>(N, D)</code>. The problem with this implementation is that <code>scatter_add_</code> requires <code>M &gt;= N</code></p>
<p>Is there any better solution? Thanks!</p>

ptrblck: <p>Would <code>scatter_add</code> with <code>dim=1</code> work?</p>
<pre><code class="lang-python">
index = torch.tensor([[0, 0, 1, 1, 0, 1],
                      [1, 1, 0, 0, 1, 0]])
data = torch.tensor([[5., 1., 7., 2., 3., 4.],
                     [5., 1., 7., 2., 3., 4.]])

torch.zeros(2, 2).scatter_add(1, index, data)
&gt; tensor([[  9.,  13.],
          [ 13.,   9.]])
</code></pre>

handesy: <p>Thanks so much for the solution!</p>

additya1998: <p>Thank you, this works. But from what I understand, this should require a lot of memory, since the source data array will be replicated <em>x</em> times if I <em>x</em> different selected sums are taken? For example, the original array is replicated thrice here:</p>
<pre><code class="lang-auto">index = torch.tensor([[0, 0, 1, 1, 0, 1],
                      [1, 1, 0, 0, 1, 0],
                      [0, 0, 1, 1, 1, 1],])

data = torch.tensor([[5., 1., 7., 2., 3., 4.],
                     [5., 1., 7., 2., 3., 4.]])
                     [5., 1., 7., 2., 3., 4.]])

torch.zeros(3, 3).scatter_add(1, index, data)
&gt; tensor([[ 9., 13.,  0.],
          [13.,  9.,  0.],
          [ 6., 16.,  0.]])
</code></pre>

ptrblck: <p>You could use <code>expand</code> to change the metadata without triggering a copy:</p>
<pre><code class="lang-python">index = torch.tensor([[0, 0, 1, 1, 0, 1],
                      [1, 1, 0, 0, 1, 0]])
data = torch.tensor([[5., 1., 7., 2., 3., 4.]])

torch.zeros(2, 2).scatter_add(1, index, data.expand(2, -1))
</code></pre>

Niels_PyTorch: <p>I have a similar question. My data is a 3x3 matrix with values (<code>src</code>), and the index contains column ids (<code>index</code>), indicating to which column each element of <code>src</code> belongs.</p>
<pre><code class="lang-auto">import torch

src = torch.tensor([[1.0, 2.0, 3.0], [2.0, 0.0, 1.0], [1.0, 3.0, 4.0]], dtype=torch.int64)

index = torch.tensor([[0, 0, 1], [0, 0, 1], [0, 0, 1]])
</code></pre>
<p>The result of the unsorted segment sum should be <code>[9, 8, 0]</code> (since there are 3 columns or segments). The sum over the first column (with id 0) is 9, the sum over the second column (with id 1) is 8, and the sum over the third column is 0 (since index does not contain any id 2). How can I achieve this with <code>scatter_add</code>?</p>