JuliaDiff / JuliaDiff/ReverseDiff.jl

Custom `rrule` not working with ReverseDiff

Open
#208 1 comment 1 reaction 0 assignees View on GitHub

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

  1. Read the whole issue, then the project's contributing guide.
  2. Comment on the issue to say you are picking it up — it saves two people doing the same work.
  3. Fork the repository and make your change on a branch.
  4. 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

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.