Recommended way of updating to new Function interface?

bamos: <p>Sorry if I missed some of the other discussions on this – In a lot of my projects I have custom PyTorch Functions defined and have been getting a lot of deprecation warning errors about updating to the staticmethod interface. <a href="https://github.com/locuslab/mpc.pytorch/blob/master/mpc/lqr_step.py#L23" rel="nofollow noopener">Here’s</a> an example of one with a lot of attributes that aren’t necessarily Tensors that I would like to access in the forward/backward pass <em>and</em> I would like to return some information from inside of the forward pass that also aren’t Tensors, which I don’t know how to do with the new interface.</p>
<p>Here’s another example that I’m thinking about updating: <a href="http://github.com/locuslab/qpth/blob/master/qpth/qp.py" rel="nofollow noopener">http://github.com/locuslab/qpth/blob/master/qpth/qp.py</a></p>
<p>And here are others that I have updated that I’ve done quite messily by passing all of the attributes into the <code>forward</code> method and putting them back into the context – a hack to get nearly the same functionality as the old interface. Is returning <code>None</code> in the backward pass for all of the additional flags the right thing to do?</p>
<ul>
<li><a href="https://github.com/oxfordcontrol/osqpth/blob/master/osqpth/osqpth.py#L39" rel="nofollow noopener">https://github.com/oxfordcontrol/osqpth/blob/master/osqpth/osqpth.py#L39</a></li>
<li><a href="https://github.com/locuslab/lml/blob/master/lml.py#L54" rel="nofollow noopener">https://github.com/locuslab/lml/blob/master/lml.py#L54</a></li>
</ul>
<p>Do you have any recommendations/thoughts before I start making more updates for this? Also how should I return auxiliary information that will never be backpropped through from a <code>Function</code> call?</p>
<p>\cc <a class="mention" href="/u/alband">@albanD</a></p>

tom: <p>Hi Brandon,</p>
<p>Hm. We should have made clear that is an old interface earlier - the “new” Function interface has been that way since PyTorch 0.2, so it is bad that you’re surprised…</p>
<aside class="quote no-group" data-username="bamos" data-post="1" data-topic="55418">
<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/bamos/48/15743_2.png" class="avatar"> bamos:</div>
<blockquote>
<p>And here are others that I have updated that I’ve done quite messily by passing all of the attributes into the <code>forward</code> method and putting them back into the context – a hack to get nearly the same functionality as the old interface. Is returning <code>None</code> in the backward pass for all of the additional flags the right thing to do?</p>
</blockquote>
</aside>
<p>I think both are the right thing to do.</p>
<p>If you’re in for a hack, you can make  <a href="https://github.com/pytorch/pytorch/pull/22582#issuecomment-511142339" rel="noopener nofollow ugc">good use of closures</a> to eliminate the administration overhead. If things go well, that will also make it easier to script-enable your functions.</p>
<p>Best regards</p>
<p>Thomas</p>

bamos: <aside class="quote no-group" data-username="tom" data-post="2" data-topic="55418">
<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/tom/48/3162_2.png" class="avatar"> tom:</div>
<blockquote>
<p>If you’re in for a hack, you can make <a href="https://github.com/pytorch/pytorch/pull/22582#issuecomment-511142339" rel="noopener nofollow ugc">good use of closures </a> to eliminate the administration overhead. If things go well, that will also make it easier to script-enable your functions.</p>
</blockquote>
</aside>
<p>Thanks! Very useful – I like that interface and may start using something like that for new functions in the future. I may take a shot at also trying to add in another level of sugar to handle non-Tensor inputs/outputs with an interface like that.</p>
<p>Also I just played around a bit more with the new Function interface and modified the ReLU example – is it ok to manually instantiate the object so I can set some attributes on it and collect some stats in it that I never need to backprop through? (Arbitrarily using numpy here to make sure it works for non-trivial objects that aren’t tensors)</p>
<pre><code class="lang-auto">import torch
import numpy as np

class MyReLU(torch.autograd.Function):
    def __init__(ctx, extra_inputs):
        ctx.extra_inputs = extra_inputs
    
    @staticmethod
    def forward(ctx, input):
        ctx.save_for_backward(input)
        print('extra_inputs:', ctx.extra_inputs)
        ctx.extra = {'some_stats': np.random.randn(10)}
        return input.clamp(min=0)

    @staticmethod
    def backward(ctx, grad_output):
        input, = ctx.saved_tensors
        grad_input = grad_output.clone()
        grad_input[input &lt; 0] = 0
        return grad_input
    
ctx = MyReLU(extra_inputs=[np.random.randn(10)])
x = torch.randn(10, requires_grad=True)
y = MyReLU.forward(ctx, x)
print(ctx.extra)
print(torch.autograd.grad(y.sum(), x))
</code></pre>

tom: <p>Note that ctx is <em>not</em> an instance of Function but rather a “backward context” object, so this won’t work. What you could do is a thing like</p>
<pre><code class="lang-python">def MyFunc(extra_inputs):
    class MyFuncFn(torch.autograd.Function):
        @staticmethod
        def foward(ctx, input):
            ...
        ...
    return MyFuncFn.apply
</code></pre>
<p>or so, but I’d have doubts whether that’s really worth it (relative to using <code>functools.partial</code> or so).<br>
I tend to use mutuable types (e.g. dictionaries) for stats.</p>
<p>Best regards</p>
<p>Thomas</p>

bamos: <p>Ah I see – in my example above my <code>torch.autograd.grad</code> call was just relying on autodiff on the <code>forward</code> method and never calling into the <code>backward</code> method.</p>
<p>And thanks! Nesting the <code>Function</code> inside a Python function is a really clean and easy way of updating my older code and not something I considered before – I updated <code>qpth</code> to do this in a few minutes and I think it’ll work well for the mpc one too with some slight modifications <img src="https://discuss.pytorch.org/images/emoji/apple/slight_smile.png?v=9" title=":slight_smile:" class="emoji" alt=":slight_smile:"></p>
<aside class="onebox githubcommit">
  <header class="source">
      <a href="https://github.com/locuslab/qpth/commit/315d5f9f6f8938e0ceade1cd63ffffe2d5f2bfb6" target="_blank" rel="nofollow noopener">github.com/locuslab/qpth</a>
  </header>
  <article class="onebox-body">
      <a href="https://github.com/bamos" target="_blank" rel="nofollow noopener">
    <img alt="bamos" src="https://avatars1.githubusercontent.com/u/707462?v=4" class="thumbnail onebox-avatar" width="90" height="90">
  </a>

<h4>
  <a href="https://github.com/locuslab/qpth/commit/315d5f9f6f8938e0ceade1cd63ffffe2d5f2bfb6" target="_blank" rel="nofollow noopener">Update to the new Function interface and update some byte-&gt;bool indexing</a>
</h4>


<div class="date">
  by <a href="https://github.com/bamos" target="_blank" rel="nofollow noopener">bamos</a>
  on <a href="https://github.com/locuslab/qpth/commit/315d5f9f6f8938e0ceade1cd63ffffe2d5f2bfb6" target="_blank" rel="nofollow noopener">03:02PM - 09 Sep 19 UTC</a>
</div>

<div class="github-commit-stats">
  changed <strong>4 files</strong>
  with <strong>169 additions</strong>
  and <strong>173 deletions</strong>.
</div>

  </article>
  <div class="onebox-metadata">
    
    
  </div>
  <div style="clear: both"></div>
</aside>