How does `scatter_add` work?

username: shinyeyes
cooked: <p>Hi,</p>
<p>I’m bit puzzled on how <code>scatter_add</code>(<a href="https://pytorch.org/docs/stable/tensors.html#torch.Tensor.scatter_add_" rel="nofollow noopener">https://pytorch.org/docs/stable/tensors.html#torch.Tensor.scatter_add_</a>) works.</p>
<h2>Problem</h2>
<p>Why can’t I reproduce the same result with for-loops?</p>
<pre><code class="lang-auto">self[index[i][j][k]][j][k] += other[i][j][k]  # if dim == 0
self[i][index[i][j][k]][k] += other[i][j][k]  # if dim == 1
self[i][j][index[i][j][k]] += other[i][j][k]  # if dim == 2
</code></pre>
<h2>Examples</h2>
<p>I’m testing the method progressively starting from a</p>
<h3>1-D tensor.</h3>
<pre><code class="lang-python">a = tf.Variable([1,2,3], dtype=tf.float32)
b = torch.tensor([1,2,3], dtype=torch.float32)

indices = np.array([0,0,1,2,1,0,2], dtype=np.int)
updates = np.array(list(range(len(indices))), dtype=np.float32)

b.scatter_add(0, torch.from_numpy(indices), torch.from_numpy(updates))
# tensor([ 7.,  8., 12.])
</code></pre>
<h3>2-D tensor.</h3>
<pre><code class="lang-python">array = np.array([
    [1,2,3],
    [4,5,6]
], dtype=np.float32)

b = torch.tensor(array, dtype=torch.float32)

indices = np.array([
    [0,0], [0,1]
], dtype=np.int)
updates = np.array(
    [[1,2], [3,4]], 
    dtype=np.float32)
</code></pre>
<p>Here I run into an error:</p>
<pre><code class="lang-python">b.scatter_add(0, torch.from_numpy(indices), torch.from_numpy(updates))

# Isn't the above the same as below?
for i in range(2):
    for j in range(2):
        b[indices[i,j]][j] += updates[i,j]
        
# tensor([[5., 4., 3.],
#         [4., 9., 6.]])
</code></pre>
<p>The following does not raise an error:</p>
<pre><code class="lang-auto">b.scatter_add(1, torch.from_numpy(indices), torch.from_numpy(updates))
# tensor([[11.,  4.,  3.],
#         [10., 17.,  6.]])


# However, the result is not as same as 
for i in range(2):
    for j in range(2):
        b[i, indices[i,j]] += updates[i,j]
# tensor([[ 8.,  4.,  3.],
#         [ 7., 13.,  6.]])
</code></pre>
<p>Thanks!</p>

username: shinyeyes
cooked: <p>Plus, is <code>scatter_add</code> compatible with <code>tf.scatter_add</code>(<a href="https://www.tensorflow.org/api_docs/python/tf/scatter_add" rel="nofollow noopener">https://www.tensorflow.org/api_docs/python/tf/scatter_add</a>)?</p>

username: ptrblck
cooked: <p>The second example yields the same result on my machine.<br>
Could you make sure to reinitialize <code>b</code> and check the result again?</p>
<p>The first call raises an error, since you have a dimension mismatch.<br>
If you expand <code>indices</code> and <code>updates</code>, the code will also return the same values:</p>
<pre><code class="lang-python">array = np.array([[1,2,3],
                  [4,5,6]],
                 dtype=np.float32)

b = torch.tensor(array, dtype=torch.float32)

indices = np.array([[0,0,0],
                    [0,1,1]],
                    dtype=np.int)
updates = np.array([[1,2,3],
                    [4,5,6]], 
                   dtype=np.float32)

print(b.scatter_add(0, torch.from_numpy(indices), torch.from_numpy(updates)))

for i in range(2):
    for j in range(3):
        b[indices[i,j]][j] += updates[i,j]
print(b)

&gt; tensor([[ 6.,  4.,  6.],
          [ 4., 10., 12.]])
&gt; tensor([[ 6.,  4.,  6.],
          [ 4., 10., 12.]])
</code></pre>

username: shinyeyes
cooked: <p>I see. So for a 2-D tensor of shape <code>(x, y)</code>, <code>y</code> has to be the same for all tensors involved in.</p>

username: shinyeyes
cooked: <p>As of PyTorch 1.1.0 and TF 1.14.0, their logic for <code>scatter_add</code> differ. So I just created a custom ops in TF for my project.</p>