Explicitly Define Gradient for Certain Computation

Say there is function F . which cannot be computed exactly but it’s gradient can be perfectly defined can I define function s.t I tell the autograd , how gradient should be computed instead of autograd figuring about by computation graph. So my requirement is as following,

function F(x) :
      #cannot be computed exactly 
      approx = some_decent_approx_of_F(x)
      return approx
function Grad_of_F(x):
      #can be computed exactly
     grad  = exact_grad_of_F(x)
     return grad

so instead of gradient computed based on computational graph of F. I want to explicitly say use Grad_of_F for computation of gradient in backward pass ? It would really helpful if someone could point to how to do it. Thanks In advance ?


You will need to use a custom autograd.Function for that. This section of the doc shows you how to do so and how to use it.

1 Like