Tf.scatter_update in pytorch

username: ZSheikhb
cooked: <p>I am wondering if anyone can refer me to the equivalent of the <code>tf.scatter_update(data, indices, batch)</code> function in pythorch?</p>

username: ptrblck
cooked: <p>Based on the <a href="https://www.tensorflow.org/api_docs/python/tf/compat/v1/scatter_update">docs</a> I think <code>torch.scatter</code> should work unless I misunderstand what the <code>_update</code> part means in TF1.</p>

username: ZSheikhb
cooked: <aside class="quote no-group" data-username="ptrblck" data-post="2" data-topic="159521">
<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/ptrblck/48/1823_2.png" class="avatar"> ptrblck:</div>
<blockquote>
<p>Based on the <a href="https://www.tensorflow.org/api_docs/python/tf/compat/v1/scatter_update" rel="noopener nofollow ugc">docs </a> I think</p>
</blockquote>
</aside>
<p>The <code>scatter_update</code> function  is given <a href="https://github.com/tensorflow/tensorflow/blob/87462bfac761435a46641ff2f10ad0b6e5414a4b/tensorflow/python/ops/state_ops.py#L383" rel="noopener nofollow ugc">here in the tensorflow repository</a>. What do you think each of inputs of  <code>scatter_update</code> function would be placed as inputs for the <code>torch.scatter</code> function?</p>

username: ptrblck
cooked: <p>I would assume <code>ref, indices, updates</code> corresponds to <code>input, (dim,) index, src</code> in <code>torch.scatter</code> where in PyTorch you would additionally specify the dimension via <code>dim</code>.</p>

username: ZSheikhb
cooked: <p>Do you think that the <code>dim</code> should be <code>zero</code> based on the definition of <code>scatter_update</code> function?</p>

username: ZSheikhb
cooked: <p>I am getting the following error message, when I use <code>torch.scatter</code> as a substitute for <code>tf.scatter_update</code>:</p>
<pre><code class="lang-auto">data: tensor([[0., 0., 0.],
        [0., 0., 0.],
        [0., 0., 0.],
        [0., 0., 0.],
        [0., 0., 0.],
        [0., 0., 0.],
        [0., 0., 0.],
        [0., 0., 0.],
        [0., 0., 0.],
        [0., 0., 0.]]), indexes:tensor([0, 1, 2]), ...src: [ 0.42924792 -0.9031867  -0.06928237]

    351         print(f"data: {data}, indexes:{indices}, ...src: {batch[:effective_batch_size]}")
--&gt; 352         data= data.scatter(dim=1,index=indices, src=torch.tensor(batch[:effective_batch_size]))
    353         print(data)
    354         #scatter_update(data, indices, batch[:effective_batch_size])

RuntimeError: Index tensor must have the same number of dimensions as self tensor
</code></pre>
<p>Any suggestion?</p>

username: ptrblck
cooked: <p>I don’t fully understand your use case so could you post the desired result tensor?<br>
If you want to assign <code>src</code> to all rows of <code>data</code> you could directly index it:</p>
<pre><code class="lang-python">data[:, idx] = src
</code></pre>
<p>or just <code>expand</code> or <code>repeat</code> <code>src</code>.</p>

username: ZSheikhb
cooked: <p>Here is the original tensorflow line and the printed results of the output</p>
<pre><code class="lang-auto">    for key in transitions._asdict().keys():
      data = getattr(self._data, key)
      batch = getattr(transitions, key)
      tf.print("data is ",data) #added prints here
      tf.print("key is ",key)
      tf.print("indices are ",indices, "effective batch size: ",effective_batch_size)
      tf.print("batch has ", batch)
      tf.scatter_update(data, indices, batch[:effective_batch_size])
      tf.print(" data after update becomes ",data)
</code></pre>
<p>The outputs:</p>
<pre><code class="lang-auto">
data is  [[0 0 0]
 [0 0 0]
 [0 0 0]
 ...
 [0 0 0]
 [0 0 0]
 [0 0 0]]
key is  s1
indices are  [0] effective batch size:  1
batch has  [[-0.989759803 0.142742947 0.788506687]]
 data after update becomes  [[-0.989759803 0.142742947 0.788506687]
 [0 0 0]
 [0 0 0]
 ...
 [0 0 0]
 [0 0 0]
 [0 0 0]]
data is  [[0 0 0]
 [0 0 0]
 [0 0 0]
 ...
 [0 0 0]
 [0 0 0]
 [0 0 0]]
key is  s2
indices are  [0] effective batch size:  1
batch has  [[-0.995679677 0.0928544 1.00487685]]
 data after update becomes  [[-0.995679677 0.0928544 1.00487685]
 [0 0 0]
 [0 0 0]
 ...
 [0 0 0]
 [0 0 0]
 [0 0 0]]
data is  [[0]
 [0]
 [0]
 ...
 [0]
 [0]
 [0]]
key is  a1
indices are  [0] effective batch size:  1
batch has  [[0.72875309]]
 data after update becomes  [[0.72875309]
 [0]
 [0]
 ...
 [0]
 [0]
 [0]]
data is  [[0]
 [0]
 [0]
 ...
 [0]
 [0]
 [0]]
key is  a2
indices are  [0] effective batch size:  1
batch has  [[0.613955]]
 data after update becomes  [[0.613955]
 [0]
 [0]
 ...
 [0]
 [0]
 [0]]
data is  [0 0 0 ... 0 0 0]
key is  discount
indices are  [0] effective batch size:  1
batch has  [1]
 data after update becomes  [1 0 0 ... 0 0 0]
data is  [0 0 0 ... 0 0 0]
key is  reward
indices are  [0] effective batch size:  1
batch has  [-9.05287075]
 data after update becomes  [-9.05287075 0 0 ... 0 0 0]
</code></pre>
<p>My code in pytorch based on your suggestion:</p>
<pre><code class="lang-auto">    for key in transitions._asdict().keys():
        data = getattr(self._data, key)
        batch = getattr(transitions, key)
        print(f"data: {data}, indexes:{indices}, ...src: {batch[:effective_batch_size]}")
        data= data.scatter(dim=1,index=indices, src=torch.tensor(batch[:effective_batch_size]))
        print(data)
</code></pre>
<p>I am wondering where I am doing a mistake here?</p>

username: ptrblck
cooked: <p>Based on your example it seems you want to use the index in <code>dim0</code> in <code>data</code> and could directly assign <code>batch</code> as seen here:</p>
<pre><code class="lang-python">data = torch.zeros(7, 3)
batch = torch.tensor([-0.989759803, 0.142742947, 0.788506687])
data[0, :] = batch
print(data)
# tensor([[-0.9898,  0.1427,  0.7885],
#         [ 0.0000,  0.0000,  0.0000],
#         [ 0.0000,  0.0000,  0.0000],
#         [ 0.0000,  0.0000,  0.0000],
#         [ 0.0000,  0.0000,  0.0000],
#         [ 0.0000,  0.0000,  0.0000],
#         [ 0.0000,  0.0000,  0.0000]])
</code></pre>
<p>Let me know if I misunderstand your examples.</p>