Inplace Fill with Index Tensor

username: btang
cooked: <p>I want to fill a data tensor with a float given some index tensor. It seems very simple and almost like <code>masked_fill_</code>, but my index tensor doesn’t have the same size as the data tensor. Can you please take a quick look at it? I am looking for something like <code>x.index_fill_(index, 42)</code> in the following example:</p>
<pre><code class="lang-auto">x = torch.Tensor([[0, 1, 2, 3, 4],
                  [ 5, 6, 7, 8, 9],
                  [10, 11, 12, 13, 14]])
index = torch.LongTensor([[1, 3], [0, 4], [2, 3]])
# After indexed fill
x = torch.Tensor([[0, 42, 2, 42, 4],
                  [42, 6, 7, 8, 42],
                  [10, 11, 42, 42, 14]])
</code></pre>

username: tom
cooked: <p>I think you want <code>scatter_</code>:</p>
<pre><code class="lang-python">x = torch.tensor([[0, 1, 2, 3, 4],
                  [ 5, 6, 7, 8, 9],
                  [10, 11, 12, 13, 14]], dtype=torch.float)
index = torch.tensor([[1, 3], [0, 4], [2, 3]])
x.scatter_(1, index, 42)
print (x == torch.tensor([[0, 42, 2, 42, 4],
                  [42, 6, 7, 8, 42],
                  [10, 11, 42, 42, 14]], dtype=torch.float))
</code></pre>
<p>For more complex tasks, you could also build an index for the first dimension and use indexing:</p>
<pre><code class="lang-python">x = torch.tensor([[0, 1, 2, 3, 4],
                  [ 5, 6, 7, 8, 9],
                  [10, 11, 12, 13, 14]], dtype=torch.float)
index = torch.tensor([[1, 3], [0, 4], [2, 3]])
index0 = torch.arange(3, dtype=torch.long)[:, None].expand(-1, 2)
x[index2.reshape(-1), index.view(-1)] = 42
print (x == torch.tensor([[0, 42, 2, 42, 4],
                  [42, 6, 7, 8, 42],
                  [10, 11, 42, 42, 14]], dtype=torch.float))
</code></pre>
<p>Best regards</p>
<p>Thomas</p>

username: shivamsaboo17
cooked: <p>In your example if x requires grad and I change the value by slicing <code>(for example x[0, :] = torch.zeros((5)))</code>, is it still differentiable? x requires grad but zeros tensor does not in my case.</p>

username: tom
cooked: <aside class="quote no-group" data-username="shivamsaboo17" data-post="3" data-topic="25439">
<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/shivamsaboo17/48/9016_2.png" class="avatar"> shivamsaboo17:</div>
<blockquote>
<p>n your example if x requires grad and I change the value by slicing <code>(for example x[0, :] = torch.zeros((5)))</code> , is it still differentiable? x requires grad but zeros tensor does not in my case.</p>
</blockquote>
</aside>
<p>The answer here is “it depends”, and it is a bit subtle.</p>
<ul>
<li>The good news is <em>yes</em>. We can do<pre><code class="lang-auto">a = torch.randn(5,5, requires_grad=True)
b = a+1
b[0] = torch.zeros(5)
b.tanh().sum().backward()
</code></pre>
and it’ll work!</li>
<li>The bad news is there are a number of caveats:
<ul>
<li>You cannot do this on a leaf tensor (e.g. <code>a</code> above) or you get an error that the leaf was pulled into the interior of the graph. So you may need to <code>.clone()</code> the leaf. That said, if you need it for leaves, doing the zeroing in <code>with torch.no_grad():</code> and manipulating <code>a.grad</code> after the backward might be a more efficient alternative.</li>
<li>This works in the above example because the original <code>b</code> is not needed to compute the backward. If you replace <code>b = a + 1</code> with <code>b = a.tanh()</code> you will get the famous <em>one of the variables needed for gradient computation has been modified by an inplace operation</em> error message because <code>a.tanh()</code> actually wants to use the result to compute the derivative (which is <code>1-result**2</code>). Again, you could do <code>b = a.tanh().clone()</code> to avoid this.</li>
</ul>
</li>
</ul>
<p>An alternative (but also with copying) is to <em>not</em> modify the things inplace that require grad, but set up a new tensor (without requires grad) and copy in the relevant bits.</p>
<pre><code class="lang-auto">b = torch.zeros_like(a)
b[1:] = a[1:]
</code></pre>
<p>This also works when copying (non-overlapping) bits from several tensors.</p>
<p>To me, the autograd handling of <em>subtensor assignments</em> is one of the most intriguing - but also subtle - things in autograd.</p>
<p>Best regards</p>
<p>Thomas</p>

username: shivamsaboo17
cooked: <p>Thanks for detailed solution…probably I understand autograd a bit better now. As my tensor was not leaf tensor it is working fine without copying</p>

username: drevicko
cooked: <p>pytorch tensors now have an <a href="https://pytorch.org/docs/stable/tensors.html#torch.Tensor.index_fill_" rel="nofollow noopener">index_fill_</a> method.</p>