ChainRulesCore.jl icon indicating copy to clipboard operation
ChainRulesCore.jl copied to clipboard

Delaration helpers

Open oxinabox opened this issue 6 years ago • 5 comments

We have @scalar_rule which basically follows the API of DiffRules.

I think we should have more.

I think @rrule which follows the API of ZygoteRules.adjoint and a matchign @ffrule

Maybe a @DoesNotExist that declares that something is nondifferentiable.

Also something like ZygoteRules.@which (probably @which_frule and @which rrule, to locate the rules that would be called.

oxinabox avatar Sep 03 '19 10:09 oxinabox

may not need @which_frule if @which frule(args...) just works already -- need to check :)

nickrobinson251 avatar Sep 03 '19 12:09 nickrobinson251

seems to work yes

julia> x, y = rand(3, 2), rand(2, 5)
([0.675862 0.272468; 0.560009 0.489604; 0.265911 0.993395], [0.194767 0.727104 … 0.881039 0.213139; 0.818924 0.0192108 … 0.897965 0.338989])

julia> z, pullback = rrule(*, x, y)
([0.354765 0.496656 … 0.840127 0.236416; 0.51002 0.41659 … 0.933037 0.28533; 0.865305 0.212429 … 1.12631 0.393426], (DNE(), Rule(getfield(ChainRules, Symbol("##57
7#579")){Array{Float64,2}}([0.194767 0.727104 0.832507 0.881039 0.213139; 0.818924 0.0192108 0.815512 0.897965 0.338989])), Rule(getfield(ChainRules, Symbol("##57
8#580")){Array{Float64,2}}([0.675862 0.272468; 0.560009 0.489604; 0.265911 0.993395]))))

julia> @which rrule(*, x, y)
rrule(::typeof(*), A::AbstractArray{#s14,2} where #s14<:Real, B::AbstractArray{#s13,2} where #s13<:Real) in ChainRules at /Users/oxinabox/JuliaEnvs/ChainRulesWorl
d/ChainRules.jl/src/rulesets/LinearAlgebra/dense.jl:75

julia> @which rrule(*, 1, 2)
rrule(::typeof(*), x::Number, y::Number) in ChainRules at /Users/oxinabox/JuliaEnvs/ChainRulesWorld/ChainRules.jl/src/rulesets/Base/base.jl:98

oxinabox avatar Sep 03 '19 13:09 oxinabox

@nickrobinson251 summerized this recently:

  • we want macros for easier writing of rules
  • we want zygote-compatible macros so it is trivial to move over the current Zygote-only rules (i.e. no code changes at all if possible, just change which macro is called)

.. another point: we need macros for some reason e.g. handling kwargs is not possible (or painful enough to basically not be possible) with a function?

oxinabox avatar Apr 20 '20 11:04 oxinabox

The rrule/frule overload also starts to get more complex when you're passing contexts around, which is necessary to handle mutation. That's a good reason to have both a rule-making macro and a something like @which_frule.

MikeInnes avatar Apr 20 '20 15:04 MikeInnes

we also need to pass an extra argument for #68

oxinabox avatar Apr 20 '20 20:04 oxinabox