JuliaDiff / JuliaDiff/ReverseDiff.jl
Custom `rrule` not working with ReverseDiff
Nobody has claimed this yet.
- Dominant language
- Julia
- Stars
- 393
- Forks
- 60
- Avg merge
- 18h 24m
- Merged PRs (30d)
- 8
Description
Hello everyone,
First off, thank you all for your work on ReverseDiff.
I'm not sure whether this issue belongs here on in the ChainRules repo, so forgive me if it's not relevant here.
I have tried defining my own pullback using ChainRules, which claims to be compatible with ReverseDiff.
However, my attempts so far have been unsuccesful.
As a minimal example, take
using ChainRules, ChainRulesCore, ReverseDiff, Plots
function new_gradient(x)
return conj((abs(x / 2) + 1)^(-2))
end
function H(a)
return (sign(a) + 1) / 2
end
function ϕ(x)
return H(x[1])
end
@scalar_rule ϕ(x) new_gradient(x)
function ChainRulesCore.rrule(::typeof(ϕ), x::AbstractArray)
y = ϕ.(x)
function array_ϕ_pullback(ȳ)
println("Test.")
ϕ̄ = NoTangent()
x̄ = @thunk new_gradient.(x) * ȳ
return ϕ̄, x̄
end
return y, array_ϕ_pullback
end
gradplot = let
xs = -5:0.01:5
vals = [ϕ(x) for x in xs]
grads = [ReverseDiff.gradient(ϕ, [x]) for x in xs]
p = plot(xs, vals, label = "ϕ")
plot!(p, xs, getindex.(grads, 1), label = "∇ϕ")
p
end
My own rrules are never called.
Have I made an elementary mistake here, or is this actually a problem with ChainRules?
Contributor guide
No contributing guide indexed for this repository
First steps
- Read the whole issue, then the project's contributing guide.
- Comment on the issue to say you are picking it up — it saves two people doing the same work.
- Fork the repository and make your change on a branch.
- Open a pull request that references the issue number.
Research direction
Start with the minimal Julia example in the issue and run the ReverseDiff.gradient call while checking whether the custom ChainRulesCore.rrule or @scalar_rule is reached. Trace the interaction between ReverseDiff and ChainRules, and consider the issue resolved when the cause of the ignored rules and the supported behavior are established.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- julia
- Domain
- devtools
- Issue type
- Bug
- Difficulty
- 4/5
- Estimated time
- 3-5 days
- Activity status
- Stale
- Clarity
- Mostly clear
- Newbie friendliness
- 35/100