Custom gradients for function. Resource for rrule and frule

KapilKhanal
Wed, 17 Jul 2024 18:32:02 +0000
General Usage

I am trying to create a custom gradients for a greens function whose term by term direct differentiation is not that accurate and we have a derivation of it’s derivatives and would like to use those. For the simple MWE example below, how can we write frules and rrules ? Also, Can anybody direct me to a good resource on Chainrules (I got confused by documentation)? I need even simpler tutorial than that.

Thank you,

function greens(m,n)
    """m,n are size (3) vectors"""
    return m' * n
end

function ∂greens(m,n)  #say this is the accurate derivative
    """Partial of greens with respect to first argument m"""
    return 2 .* m .* n
end

f = greens([1,2,3], [3,2,0]) #7
∂f_∂m =  ∂greens([1,2,3], [3,2,0]) #[6,8,0]
# print("AD gradients \n ")
print(f)
print(" \n and \n")
print(∂f_∂m)

19 posts - 3 participants