Disable in-place correctness version check / any other workaround

username: patrick-kidger
cooked: <p>I realise the title sounds dangerous.</p>
<p>I’d like to do something akin to the following:</p>
<pre><code class="lang-auto">a = torch.rand(1, requires_grad=True)
b = torch.rand(2, requires_grad=True)
c = torch.rand(3, requires_grad=True)
x = torch.rand(1, requires_grad=True)
y = torch.empty(4)
y[0] = x
y[1] = y[:1].dot(a)
y[2] = y[:2].dot(b)
y[3] = y[:3].dot(c)
y.sum().backward()
</code></pre>
<p>If <code>a</code>, <code>b</code>, <code>c</code> don’t require gradients then this works fine. However if they do then an error is thrown, because <code>y</code> is needed to calculate the gradients wrt <code>a</code>, <code>b</code>, <code>c</code>, and it’s been modified in-place.</p>
<p>Except - it hasn’t! The only bits that have been modified are the bits that don’t affect the gradient calculation. Not that I expect the correctness checks to be able to detect that level of detail.</p>
<p>The obvious workaround is to split up <code>y</code> into lots of little tensors, but that would make each dot product slower, as we’d have to drop back into Python to do the dot product. This is inside a hot loop that I really want to be as optimised as possible.</p>
<p>Are there any clever workarounds that I might be able to pull, or any way I can turn off the hand-holding of the version tracker? (I already tried modifying the version after each assignment but apparently it’s read-only.)</p>

username: albanD
cooked: <p>Hi,</p>
<p>Yes this is an annoying limit of the inplace detection.<br>
Going around the version counter is the only case where we still have to use <code>.data</code>. Unfortunately, <code>.data</code> has also the other effect of detaching the Tensor and so no gradient will be able to flow back through that computation.</p>
<p>The only way around that I can see is to wrap it into a custom Function. Where you use <code>.data</code> in the forward to hide the version changes.<br>
But that means you have to write a custom backward for that <img src="https://discuss.pytorch.org/images/emoji/apple/confused.png?v=9" title=":confused:" class="emoji" alt=":confused:"></p>
<p>Note that a cleaner API to stop sharing version while preserving the autograd is in our todo list for the <code>.data</code> deprecation task here: <a href="https://github.com/pytorch/pytorch/issues/30987">https://github.com/pytorch/pytorch/issues/30987</a><br>
But we didn’t had time to get around to do it yet.</p>

username: tom
cooked: <p>Just because doing taxes and accounting makes me want to have a little <em>who knows the most evil autograd tricks</em> competition with <a class="mention" href="/u/alband">@albanD</a> and <a class="mention" href="/u/ptrblck">@ptrblck</a>, here is my  entry:</p>
<pre><code class="lang-python">class DangerousDot(torch.autograd.Function):
    @staticmethod
    def forward(ctx, a, b):
        ctx._a = a
        ctx._b = b
        return a.dot(b)
    @staticmethod
    def backward(ctx, grad_out):
        return grad_out * ctx._b, grad_out * ctx._a


a = torch.randn(10, requires_grad=True, dtype=torch.double)
b = torch.randn(10, requires_grad=True, dtype=torch.double)
torch.autograd.gradcheck(DangerousDot.apply, (a, b))

a = torch.rand(1, requires_grad=True)
b = torch.rand(2, requires_grad=True)
c = torch.rand(3, requires_grad=True)
x = torch.rand(1, requires_grad=True)
y = torch.empty(4)
y[0] = x
y[1] = DangerousDot.apply(y[:1], a)
y[2] = DangerousDot.apply(y[:2], b)
y[3] = DangerousDot.apply(y[:3], c)
y.sum().backward()
</code></pre>
<p><strong>Explanation:</strong> The bit that is responsible for the version check in <code>autograd.Function</code>s is <code>ctx.save_for_backward(...)</code> and <code>ctx.saved_tensors</code>. So avoiding that for <code>a</code> and <code>b</code> and assigning them to the ctx directly bypasses the checks.</p>
<p>Strictly don’t use this at home, it’s evil and dangerous!</p>
<p>Best regards</p>
<p>Thomas</p>

username: albanD
cooked: <blockquote>
<p><em>who knows the most evil autograd tricks</em> competition</p>
</blockquote>
<p><img src="https://discuss.pytorch.org/images/emoji/apple/smiley.png?v=12" title=":smiley:" class="emoji only-emoji" alt=":smiley:" loading="lazy" width="20" height="20"></p>
<p>Note as well that tom’s code will create undestructible ref cycle and will leak memory on top of being very wrong <img src="https://discuss.pytorch.org/images/emoji/apple/smiley.png?v=12" title=":smiley:" class="emoji" alt=":smiley:" loading="lazy" width="20" height="20"><br>
The original purpose of <code>save_for_backward()</code> is actually to avoid this cycle <img src="https://discuss.pytorch.org/images/emoji/apple/wink.png?v=12" title=":wink:" class="emoji" alt=":wink:" loading="lazy" width="20" height="20"></p>
<p>I would suggest this NotSoDangerousDot:</p>
<pre data-code-wrap="python"><code class="lang-python">class NotSoDangerousDot(torch.autograd.Function):
    @staticmethod
    def forward(ctx, a, b):
        prod = a * b
        ctx.save_for_backward(b)
        ctx.prod = prod
        return prod.sum()

    @staticmethod
    def backward(ctx, grad_out):
        b, = ctx.saved_tensors
        prod = ctx.prod
        return grad_out * b, grad_out * (prod / b)
</code></pre>
<p>Which might be slightly slower but should be correct in all ways! (cheating the autograd version checker is bad! haha)</p>

username: patrick-kidger
cooked: <p>Haha, thanks both! I ended up going down a very similar line to what you’ve just described, so here’s my entry to the competition:</p>
<pre><code class="lang-auto">class _AssignInAVerySafeManner(torch.autograd.Function):
    @staticmethod
    def forward(ctx, scratch, value, index):
        ctx.index = index
        scratch.data[index] = value
        return scratch

    @staticmethod
    def backward(ctx, grad_scratch):
        return grad_scratch, grad_scratch[ctx.index], None


a = torch.rand(1, requires_grad=True)
b = torch.rand(2, requires_grad=True)
c = torch.rand(3, requires_grad=True)
x = torch.rand((), requires_grad=True)
y = torch.empty(4)

y = _AssignInAVerySafeManner.apply(y, x, 0)
for i, letter in zip((1, 2, 3), (a, b, c)):
    y_ = y[:i].dot(letter)
    y = _AssignInAVerySafeManner.apply(y, y_, i)
y.sum().backward()
</code></pre>
<p>This does use the to-be-deprecated <code>.data</code>. though. Do you have any idea roughly when that’s likely to disappear?</p>

username: tom
cooked: <aside class="quote no-group" data-username="albanD" data-post="4" data-topic="90738">
<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/alband/48/215_2.png" class="avatar"> albanD:</div>
<blockquote>
<p>Note as well that tom’s code will create undestructible ref cycle and will leak memory on top of being very wrong <img src="https://discuss.pytorch.org/images/emoji/apple/smiley.png?v=12" title=":smiley:" class="emoji" alt=":smiley:" loading="lazy" width="20" height="20"><br>
The original purpose of <code>save_for_backward()</code> is actually to avoid this cycle <img src="https://discuss.pytorch.org/images/emoji/apple/wink.png?v=12" title=":wink:" class="emoji" alt=":wink:" loading="lazy" width="20" height="20"></p>
</blockquote>
</aside>
<p>Touché! But you have NaNs if b is 0.</p>
<p>More seriously, I’d probably throw cloning.</p>
<p>Best regards</p>
<p>Thomas</p>