Custom ChainRulesCore rrule with ReverseDiff

By default, ReverseDiff ignores all rrules, and one must opt into each rule one wants. Unfortunately this isn’t documented, but there’s an open PR to do so.

Here’s how to opt in:

julia> ReverseDiff.gradient(f, ones(3))
3-element Vector{Float64}:
 2.0
 2.0
 2.0

julia> ReverseDiff.@grad_from_chainrules f(x::TrackedArray)

julia> ReverseDiff.gradient(f, ones(3))
"Calling custom pullback" = "Calling custom pullback"
3-element Vector{Float64}:
 2.0
 2.0
 2.0
4 Likes