Opting out of rules

[mzgubic]

The conclusion of the discussion spread across #343, JuliaDiff/ChainRules.jl#337, #347, and JuliaDiff/ChainRules.jl#232 (and possibly elsewhere), has been to prefer writing abstractly typed rules.

Two shortcomings of this approach are that

1. In some cases, such as when the primal function is specialised for a subtype (e.g. Diagonal<:AbstractArray), the fallback rrule returning a Matrix is not correct. Solving that is the topic of Projecting Cotangents #286
2. In other cases, using AD to differentiate through the specialised method is faster than falling back on the generic rule. For these, we need a way to opt out of using the generic fallback rule, and let AD compute the backward pass. Solving that is the topic of this issue.

There are a number of possible ways we could facilitate opting-out:

• Option A) Return nothing. This has been tried in Take nothing seriously FluxML/Zygote.jl#967 but advised against in Use ChainRules RuleConfig FluxML/Zygote.jl#990 (comment) for Zygote. According to @oxinabox this would work in Diffractor, but not in Nabla? (please just edit this if that's wrong)
• Option B) Have a no_rrule function that would be used to designate a particular signature does not have a rule.
• Option C) Have a @no_rrule macro which would take a signature and generate code which returns nothing (as above), and automatically add the method to an internal list, which would be accessed by Zygote and Nabla.

Have I missed another option we've discussed?

[oxinabox]

Option A-2)
The first version of the ChainRules support in Zygote was Zygote#291.
This checked the return from rrule at runtime, rather than looking up the MethodTable in the generated function.
While I say run-time, it should have compiled that branch out of existance by type-inferring whether or not Nothing would be returned, and thus specializing away the check (til naturally invalidated by a rule being added).
A problem with it was that Zygote became super-slow, because it didn't type-infer and specialize away the branch.
because Zygote and type-inference are not friends.
But I think it might work if rather than using a if it used dispatch on the return value from rrule.
That might specialize properly, into a static dispatch.

cc @Keno does any of the options above work worse for Diffractor than current setup?
From some past discussions my impression is returning nothing is good for Diffractor, right?

[Keno]

All of the mentioned options have the same problem, which is that you're trying to reflect on global state inside a generated function, which is not allowed. You're just moving the state around, but they're all unsound in exactly the same way and will cause similar issues. Diffractor does not do any such reflection (in the generated function version at least), but instead is just careful about inference. I think returning `nothing` is still probably the cleanest way to do it. An alternative would be to do the method table query outside the generated function, but I don't see why that should be any easier on the compiler than just calling it and checking for nothing.

…

On Mon, Jun 21, 2021, 09:06 Lyndon White ***@***.***> wrote: The first version of the ChainRules support in Zygote was Zygote#291 <https://github.com/FluxML/Zygote.jl/pull/291/files#diff-7511b224d7f3ebb56465690de8e307422e3c9798a22bdd4e960d5c86ba6528aaR32-R42> . This checked the return from rrule at runtime, rather than looking up the MethodTable in the generated function. While I say run-time, it should have compiled that branch out of existance by type-inferring whether or not Nothing would be returned, and thus specializing away the check (til naturally invalidated by a rule being added). A problem with it was that Zygote became super-slow, because it *didn't* type-infer and specialize away the branch. because Zygote and type-inference are not friends. But I think it might work if rather than using a if it used dispatch on the return value from rrule. That might specialize properly, into a static dispatch. cc @Keno <https://github.com/Keno> does any of the options above work worse for Diffractor than current setup? From some past discussions my impression is returning nothing is good for Diffractor, right? — You are receiving this because you were mentioned. Reply to this email directly, view it on GitHub <#377 (comment)>, or unsubscribe <https://github.com/notifications/unsubscribe-auth/AAJ3LF6JBQAOW2L36B43GLLTT42KXANCNFSM47BQEZJA> .

[oxinabox]

So Diffractor does something like my Option A-2), I guess?

[Keno]

Yes, that's right.

…

On Mon, Jun 21, 2021, 09:25 Lyndon White ***@***.***> wrote: So Diffractor does something like my Option A-2), I guess? — You are receiving this because you were mentioned. Reply to this email directly, view it on GitHub <#377 (comment)>, or unsubscribe <https://github.com/notifications/unsubscribe-auth/AAJ3LF4RLYLDMX6M44U4NCTTT44THANCNFSM47BQEZJA> .