-
Notifications
You must be signed in to change notification settings - Fork 60
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
rules for specific ADs #270
Comments
I like the approach. I have started adding comments next to new rules explaining why they are added in the hopes of something like this in the future. Rather than hardcoding the specific AD's in the spirit of ChainRules, perhaps it makes more sense to try to encode in the signature of the rule itself the reason why it is added, perhaps using a traits-based approach or a type union. Then some hot new AD doesn't depend on us recognizing it and modifying our rules. Rather, when hooking into ChainRules it could opt in to specific categories of rules. |
Yeah traits could be a cool way to do this. Something like: configured_rrule = ConfiguredRRule(
NoMutation, NoInplaceAccumulation, HatesLoops;
forward_ad=Zygote._pushforward, #Or maybe just `nothing` as not provided.
reverse_ad=Zygote._pullback,
) for Nabla: configured_rrule = ConfiguredRRule(
NoMutation, InplaceAccumulation, IsOkWithLoops;
forward_ad=Nabla.fmad,
reverse_ad=Nabla.∇,
) |
Traits make sense -- provided we're able to provide an escape hatch so that AD implementers can always express "no, I really do want to use only that very particular rule.". |
Closed by #363 |
Lets assume that Diffractor is going to be better at some things than Zygote is and, as a consequence, there exist rules that we don't want to Diffractor to hit (since it generates perfectly good code anyway) but we do want Zygote to hit. Presently, we don't have a way to specify this.
It would be simple to achieve this via dispatch, and adding an additional argument to ones
rrule
s, making them have the following signature:This would also alleviate some of our existing headaches surrounding
rrule
s for "very abstractly typed" arguments, since we could implement generic versions of things that we're sure ought to work (e.g.*(::Matrix{Float64}, Matrix{Float64})
) inChainRules
, without requiring package authors to compromise on existing choices that they've made -- for exampleZygote
uses very abstract types for lots of things and, while I don't like it, it would be really very breaking to make changes to it at this point in time.Recall that there are essentially 3 reasons (I think?) to implement a rule:
Rules of type 1 are those for which you would consider writing a very generic
rrule
, so you would probably write them to accept anyAbstractAD
.Rules of type 2 are somewhat borderline and would probably need to be done on a case-by-case basis. For example, you might write a custom adjoint for a function involving a for-loop if using
Zygote
, but might not need the rule at all if usingDiffractor
. WhileZygote
can usually differentiate through for-loops, it tends to be slow.Rules of type 3 are prime candidates for AD-specific rules, since different ADs are able to differentiate through different language features.
This is related to #68 in that we're talking about including some kind of additional information about what ADs to use, but the underlying problem that it addresses is somewhat different.
The text was updated successfully, but these errors were encountered: