Not sure I got you. I tried to follow that guide: Extending PyTorch — PyTorch 1.8.0 documentation 1
Do you suggest grad_output * grad_s ?
So, I’ve tried:
@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

still get:
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)