Gradcheck fails despite correct implementation in backward

ichernob: <p>Fails for the following function:</p>
<pre><code class="lang-auto">class FRUFunction(Function):
    @staticmethod
    def forward(ctx, s, a):
        k = 1 / (a + s - a * s)

        ctx.save_for_backward(s, k, a)
        
        return k * s

    @staticmethod
    def backward(ctx, grad_output):
        s, k, a = ctx.saved_tensors
        grad_s = grad_a = None
        
        if ctx.needs_input_grad[0]:
            grad_s = a / (a + s - a*s)** 2

        return grad_s, grad_a
</code></pre>
<p>Check with code:</p>
<pre><code class="lang-auto">input = torch.tensor([[0.0971, 0.5413, 0.8107]], dtype=torch.double, requires_grad=True)
alpha = torch.tensor([[0.2904, -0.3183, 0.7001]], dtype=torch.double, requires_grad=False)
fru_input = (
    input,
    alpha,
)
test = gradcheck(fru, fru_input, eps=1e-3)
</code></pre>
<p>It throws runtime exception:</p>
<pre><code class="lang-auto">RuntimeError: Jacobian mismatch for output 0 with respect to input 0,
numerical:tensor([[ 2.2495,  0.0000,  0.0000],
        [ 0.0000, -2.0370,  0.0000],
        [ 0.0000,  0.0000,  0.7869]], dtype=torch.float64)
analytical:tensor([[ 2.2495,  2.2495,  2.2495],
        [-2.0370, -2.0370, -2.0370],
        [ 0.7869,  0.7869,  0.7869]], dtype=torch.float64)
</code></pre>
<p>NOTE: It even fails for</p>
<pre><code class="lang-auto">@staticmethod
def forward(ctx, s, a):
    k = 1 / (a + s - a * s)
    ctx.save_for_backward(s, k, a)
    
    return k

@staticmethod
def backward(ctx, grad_output):
    s, k, a = ctx.saved_tensors
    grad_s = grad_a = None
    if ctx.needs_input_grad[0]:
        grad_s = (a-1) / ((a + s - a*s) ** 2)
    
    return grad_s, grad_a
</code></pre>
<p>and for</p>
<pre><code class="lang-auto">@staticmethod
def forward(ctx, s, a):
    k = a + s - a * s
    ctx.save_for_backward(s, k, a)
    
    return k

@staticmethod
def backward(ctx, grad_output):
    s, k, a = ctx.saved_tensors
    grad_s = grad_a = None
    if ctx.needs_input_grad[0]:
        grad_s = a-1
    
    return grad_s, grad_a
</code></pre>
<p>and for</p>
<pre><code class="lang-auto">@staticmethod
def forward(ctx, s, a):
    k = a * s
    ctx.save_for_backward(s, k, a)
    
    return k

@staticmethod
def backward(ctx, grad_output):
    s, k, a = ctx.saved_tensors
    grad_s = grad_a = None
    if ctx.needs_input_grad[0]:
        grad_s = a
    
    return grad_s, grad_a
</code></pre>

albanD: <p>Hey,</p>
<p>I looks like none of your formulas is using the <code>grad_output</code> parameter?</p>
<p>What the backward function for a functin f should be computing is the vector Jacobian product between <code>grad_output</code> and the Jacobian of f.</p>

ichernob: <p>Not sure I got you. I tried to follow that guide: <a href="https://pytorch.org/docs/stable/notes/extending.html?highlight=gradcheck" class="inline-onebox" rel="noopener nofollow ugc">Extending PyTorch — PyTorch 1.8.0 documentation</a></p>
<p>Do you suggest <code>grad_output * grad_s</code> ?</p>
<p>So, I’ve tried:</p>
<pre><code class="lang-auto">@staticmethod
    def forward(ctx, s, a):
        k = 1 / (a + s - a * s)

        ctx.save_for_backward(s, a, k)

        return k

@staticmethod
def backward(ctx, grad_output):
    s, a, k = ctx.saved_tensors
    grad_s = None
    grad_s = (1 - a) * (k ** 2)

    return grad_output * grad_s, None
</code></pre>
<p>still get:</p>
<pre><code class="lang-auto">RuntimeError: Jacobian mismatch for output 0 with respect to input 0,
numerical:tensor([[-5.4966,  0.0000,  0.0000],
        [ 0.0000, -8.4367,  0.0000],
        [ 0.0000,  0.0000, -0.3371]], dtype=torch.float64)
analytical:tensor([[5.4966, 0.0000, 0.0000],
        [0.0000, 8.4366, 0.0000],
        [0.0000, 0.0000, 0.3371]], dtype=torch.float64)
</code></pre>

ichernob: <p>Well, it seems that piece does work for me:</p>
<pre><code class="lang-auto">    @staticmethod
    def forward(ctx, s, a):
        k = 1 / (a + s - a * s)

        ctx.save_for_backward(s, a, k)

        return k * s

    @staticmethod
    def backward(ctx, grad_output):
        s, a, k = ctx.saved_tensors
        grad_s = None
        grad_s = k + s * (a - 1) * (k ** 2)

        return grad_output * grad_s, None
</code></pre>

albanD: <p>That looks better yes.</p>
<p>In the tutorial you linked, the formula is implemented properly there and uses grad_output.</p>