How to pass intermediate value in forward function into self class instantiation?

HaFred: <p>My problem is best illustrated with an example.<br>
I wrapped the <code>autograd.Function</code> class into the <code>uniform</code> function and used it in the forward function of <code>target_fn</code> class, which works fine as below.</p>
<pre><code class="lang-auto">def uniform(k):
  class qfn(torch.autograd.Function):

    @staticmethod
    def forward(ctx, input):
      ...
      return out

    @staticmethod
    def backward(ctx, grad_output):
      ...
      return grad_input
  
  return qfn().apply
  
class target_fn(nn.Module):
  def __init__(self, bit):
    super(target_fn, self).__init__()
    ...
    self.uni = uniform(k=bit)  

  def forward(self, x):
#     self.max = max(x)
    output = self.uni(x)
    return output  
</code></pre>
<p>However, I want the <code>self.max</code> computed in the <code>target_fn.forward</code> to be utilized by the <code>uniform</code> for some preprocessing. And I tried something like below but error prompted:</p>
<pre><code class="lang-auto">def uniform(k, max_value):
  class qfn(torch.autograd.Function):

    @staticmethod
    def forward(ctx, input):
      ...
      out = out / max_value
      return out

    @staticmethod
    def backward(ctx, grad_output):
      ...
      return grad_input
    
  return qfn().apply
  
class target_fn(nn.Module):
  def __init__(self, bit):
    super(target_fn, self).__init__()
    ...
    self.uni = uniform(k=bit, max_value = self.max)   # error: target_fn instantiation does not have self.max

  def forward(self, x):
    self.max = max(x)
    output = self.uni(x)
    return output
</code></pre>
<p>Funny thing is that I can actually access the <code>target_fn.max</code> in debugger, but the line <code>self.uni = uniform(k=bit, max_value = self.max)</code> cannot.<br>
How could I fix it? Thank you.</p>

albanD: <p>Hi,</p>
<p>Why not just pass the <code>self.max</code> as a second argument to your <code>qfn</code> custom Function?<br>
If it does not require gradients, you can simply return <code>None</code> in the backward pass for it.</p>

HaFred: <p>Hi Alban,</p>
<p>Thank you for your reply. Sorry that I don’t quite get what you meant. As you see, in the second snippet, I have passed the <code>self.max</code> as the second argument of the uniform function. Could you elaborate on your suggestion?</p>

albanD: <p>I meant the <code>qfn</code> function so that the forward is <code>def forward(ctx, input, max_value)</code>.</p>

HaFred: <p>Hi,<br>
Actually the <code>self.max</code> is determined in the <code>target_fn.forward</code>. I think my question could be better rephrased like this:<br>
How can we access the <code>self.max</code> in the <code>target.__init__</code> to get through this error? I know register_forward_hook could help on the activation interception in <code>forward</code>, I am just curious if we need to pass the <code>self.max</code> as an argument in the <code>uniform()</code> in <code>target_fn.__init__</code>, what should I do?</p>

albanD: <p>I don’t understand why you want <code>max</code> to be saved on the Module?<br>
And why you cannot just pass it along?</p>
<pre><code class="lang-python">def uniform(k):
  class qfn(torch.autograd.Function):

    @staticmethod
    def forward(ctx, input, max_value):
      ...
      out = out / max_value
      return out

    @staticmethod
    def backward(ctx, grad_output):
      ...
      EDIT: need to return None for non-differentiable inputs
      return grad_input, None
  
  return qfn().apply

class target_fn(nn.Module):
  def __init__(self, bit):
    super(target_fn, self).__init__()
    ...
    self.uni = uniform(k=bit)

  def forward(self, x):
    max = max(x)
    output = self.uni(x, max)
    return output
</code></pre>

HaFred: <p>I don’t know if pass into <code>forward</code> function is a good idea? I tried your solution but prompted <code>RuntimeError: function qfnBackward returned an incorrect number of gradients (expected 2, got 1)</code>. I think the <code>backward</code> might consider the <code>max_value</code> for backprop also, so I tried to set the <code>max_value.requires_grad = False</code>. But the error still remained.</p>
<pre><code class="lang-auto">def uniform(k):
  class qfn(torch.autograd.Function):

    @staticmethod
    def forward(ctx, input, max_value):
      ...
      out = out / max_value
      return out

    @staticmethod
    def backward(ctx, grad_output):
      ...
      return grad_input
  
  return qfn().apply

class target_fn(nn.Module):
  def __init__(self, bit):
    super(target_fn, self).__init__()
    ...
    self.uni = uniform(k=bit)

  def forward(self, x):
    max = max(x)
    max.requires_grad = False
    output = self.uni(x, max)
    return output
</code></pre>
<p>The reason I want <code>max</code> to be saved on the module is to properly construct the <code>uniform</code>. You see, the <code>uniform.qfn.forward</code> contains some logic based on the <code>max</code>.<br>
Thanks again, albanD.</p>

albanD: <p>Ho sorry I forgot to update the backward in my code sample above <img src="https://discuss.pytorch.org/images/emoji/apple/confused.png?v=9" title=":confused:" class="emoji" alt=":confused:">  Done now.<br>
Any input that is not a Tensor that requires grad (either not a Tensor or a Tensor that does not require gradients) should have the backward return None for them.</p>

HaFred: <p>Hi, I did the <code>return grad_input, None</code> but the error still occurred.<br>
<code>RuntimeError: function qfnBackward returned an incorrect number of gradients (expected 2, got 1)</code><br>
Could you reproduce this error on your side? If the <code>forward</code> has 2 arguments excluding the context ctx, and if the backward really gonna work with <code>grad_input, None</code>? Thanks again.</p>

albanD: <p>Yes that does work on my end. You should double check your code to make sure you don’t have an old version of the custom Function around.</p>