Using autograd.Function to compute gradient for a function

arunppsg: <p>I am learning <code>autograd</code> and I was following the tutorials <a href="https://pytorch.org/docs/stable/notes/extending.html" rel="noopener nofollow ugc">extending pytorch</a>.</p>
<p>I would like to compute gradients wrt x and y for the following function.</p>
<pre><code class="lang-auto">import torch
from torch.autograd import Function

class Func(Function):
    
    @staticmethod
    def forward(x, v):
        a = torch.tanh(x)
        o = torch.matmul(a, v)
        return o, a
    
    @staticmethod
    def setup_context(ctx, inputs, outputs):
        x, v = inputs
        o, a = outputs
        ctx.save_for_backward(x, v, a)
        
    @staticmethod
    def backward(ctx, grad_o, grad_a):
        x, v, a = ctx.saved_tensors

        dL_da = torch.matmul(grad_o, v.transpose(1, 0))
        dL_dv = torch.matmul(a.transpose(1, 0), grad_o)
        
        dL_dx = (dL_da / ((torch.cosh(x)) ** 2))
        return dL_dx, dL_dv
</code></pre>
<p>When I run <code>gradcheck</code> as</p>
<pre><code class="lang-auto">x = torch.randn(1, 3)

func = Func.apply
x = torch.randn(2, 3, requires_grad=True, dtype=torch.double)
y = torch.randn(3, 2, requires_grad=True, dtype=torch.double)
test = gradcheck(func, (x, y), eps=1e-6, atol=1e-4)
</code></pre>
<p>it throws me the following traceback:</p>
<pre><code class="lang-auto">GradcheckError: Jacobian mismatch for output 1 with respect to input 0,
numerical:tensor([[0.4000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000],
        [0.0000, 0.7561, 0.0000, 0.0000, 0.0000, 0.0000],
        [0.0000, 0.0000, 0.9592, 0.0000, 0.0000, 0.0000],
        [0.0000, 0.0000, 0.0000, 0.9947, 0.0000, 0.0000],
        [0.0000, 0.0000, 0.0000, 0.0000, 0.9674, 0.0000],
        [0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.4168]], dtype=torch.float64)
analytical: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.],
        [0., 0., 0., 0., 0., 0.]], dtype=torch.float64)
</code></pre>
<p>I am not sure what I am doing wrong here. Can someone help me figure out?</p>
<h3>
<a name="other-notes-1" class="anchor" href="#other-notes-1"></a>Other notes</h3>
<p>I tried out simple functions - <code>tanh</code> and <code>matmul</code> and they seem to work fine:</p>
<pre><code class="lang-auto">class Mm(Function):
    
    @staticmethod
    def forward(x, y):
        return torch.matmul(x, y)
    
    @staticmethod
    def setup_context(ctx, inputs, output):
        x, y = inputs
        ctx.save_for_backward(x, y)
        
    @staticmethod
    def backward(ctx, grad_output):
        x, y = ctx.saved_tensors
        dx = torch.matmul(grad_output, y.transpose(1, 0))
        dy = torch.matmul(x.transpose(1, 0), grad_output)
        return dx, dy

class Tanh(Function):
    
    @staticmethod
    def forward(x):
        return torch.tanh(x)
    
    @staticmethod
    def setup_context(ctx, inputs, output):
        x, = inputs
        ctx.save_for_backward(x)
        
    @staticmethod
    def backward(ctx, grad_y):
        x, = ctx.saved_tensors
        return grad_y * (1 / ((torch.cosh(x)) ** 2))
    
mm = Mm.apply
x = torch.randn(2, 3, requires_grad=True, dtype=torch.double)
y = torch.randn(3, 2, requires_grad=True, dtype=torch.double)
test = gradcheck(mm, (x, y), eps=1e-6, atol=1e-4)

tanh = Tanh.apply
x = torch.randn(1, 3, requires_grad=True, dtype=torch.double)
test = gradcheck(tanh, x, eps=1e-6, atol=1e-4)
</code></pre>

AlphaBetaGamma96: <aside class="quote no-group" data-username="arunppsg" data-post="1" data-topic="178979">
<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/arunppsg/48/52196_2.png" class="avatar"> arunppsg:</div>
<blockquote>
<pre><code class="lang-auto">func = Func.apply
</code></pre>
</blockquote>
</aside>
<p>Shouldn’t the <code>.apply</code> method contain arguments? That’s probably why you get a zero tensor for your analytical results. You should probably defined <code>func</code> as,</p>
<pre><code class="lang-auto">def func(x, v):
    return Func.apply(x, v)
</code></pre>

arunppsg: <p>The <a href="https://pytorch.org/docs/stable/notes/extending.html#example" rel="noopener nofollow ugc">doc</a> says we can do it either way:</p>
<pre><code class="lang-auto"># Option 1: alias
linear = LinearFunction.apply

# Option 2: wrap in a function, to support default args and keyword args.
def linear(input, weight, bias=None):
    return LinearFunction.apply(input, weight, bias)
</code></pre>

AlphaBetaGamma96: <aside class="quote no-group" data-username="arunppsg" data-post="3" data-topic="178979">
<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/arunppsg/48/52196_2.png" class="avatar"> arunppsg:</div>
<blockquote>
<p>The <a href="https://pytorch.org/docs/stable/notes/extending.html#example" rel="noopener nofollow ugc">doc</a> says we can do it either way:</p>
<pre><code class="lang-auto">
</code></pre>
</blockquote>
</aside>
<p>Fair enough.</p>
<p>I’ve just realized you haven’t used <code>grad_a</code> in your code, could that might be why you get a zero tensor?</p>

arunppsg: <p>That’s another question which I have - when I can compute <code>grad_a</code> from <code>grad_o</code>, should I use either the computed value or  <code>grad_a</code> from output. If I am not wrong, <code>dL_da</code> should be equal to <code>grad_a</code>.</p>
<p>I tried the following variant using <code>grad_a</code> but still getting same error.</p>
<pre><code class="lang-auto">import torch
from torch.autograd import gradcheck
from torch.autograd import Function

class Func(Function):
    
    @staticmethod
    def forward(x, v):
        a = torch.tanh(x)
        o = torch.matmul(a, v)
        return o, a

    @staticmethod
    def setup_context(ctx, inputs, outputs):
        x, v = inputs
        o, a = outputs
        ctx.save_for_backward(x, v, a)
        
    @staticmethod
    def backward(ctx, grad_o, grad_a):
        x, v, a = ctx.saved_tensors

        dL_dv = torch.matmul(a.transpose(1, 0), grad_o)
        
        dL_dx = (grad_a / ((torch.cosh(x)) ** 2))
        return dL_dx, dL_dv


x = torch.randn(1, 3)

func = Func.apply
x = torch.randn(2, 3, requires_grad=True, dtype=torch.double)
y = torch.randn(3, 2, requires_grad=True, dtype=torch.double)
test = gradcheck(func, (x, y), eps=1e-6, atol=1e-4)
</code></pre>

AlphaBetaGamma96: <p>Are you sure your backward formula is correct? There should be a (1 - tanh(x)**2) term, as that’s the derivative of tanh(x) w.r.t x.</p>
<p>The matrix cookbook should be able to help with the linear algebra.</p>

soulitzer: <pre><code class="lang-auto">
    def backward(ctx, grad_o, grad_a):
        x, v, a = ctx.saved_tensors

        dL_da = grad_a + torch.matmul(grad_o, v.transpose(1, 0))
        dL_dv = torch.matmul(a.transpose(1, 0), grad_o)

        dL_dx = dL_da * (1 / (torch.cosh(x) ** 2))
        return dL_dx, dL_dv
</code></pre>
<p>AlphaBetaGamma has the right idea, you missed the grad_a. Technically though, your backward isn’t wrong if you never use <code>a</code> in any gradient computation. I would guess that it is only returned here so you could save it for backward (usually you would then have a wrapper function that filters out the <code>a</code> from the outputs).</p>
<p>The formula is detected to be wrong by gradcheck because it try to backprop from <code>a</code> back to the inputs, and since <code>a</code> is indeed a differentiable function of the inputs, the numerical and analytical gradients will differ</p>