Support frules with keyword arguments #266
Conversation
| # then check if the frule for f accepts keywords | ||
| # This silently discards tangents of the kw-args | ||
| # TODO: should we error if they nonzero? | ||
| r = _frule(args_partials, args_primals...; primal(kw)...) |
There was a problem hiding this comment.
I'm not sure I like this fallback. The problem is that if we support it at first order, we need to have the same logic at higher orders and it just becomes complicated.
There was a problem hiding this comment.
If we don't have this, then we do not need this PR at all, since writing the rule for kwcall currently does work.
This particular logic is the code of this PR.
It's what makes it possible to write frules that have keywords args that match to functions with keyword args.
I don't understand what the problem with this at higher order is?
| end | ||
|
|
||
| _frule(partials, primals...; kwargs...) = frule(DiffractorRuleConfig(), partials, primals...; kwargs...) | ||
| function _frule(::NTuple{<:Any, AbstractZero}, f, primal_args...; kwargs...) |
There was a problem hiding this comment.
I suspect for performance, you probably want to define the non-kwargs version separately, so that higher-order AD doesn't have to AD through all the kwargs dispatch logic, like we do in ChainRules.
Closes #244
I need to write an frule for
rangeat some point.Questions:
Do we need to support writing rules for
Core.kw? right now (without this PR) it is the only way we support frules with keywords. ANd even with adding frule support if one needs to carry deriviatives though keywords it is currently the only way (and likely to remain the only way til a bigger API overhaul in ChainRules 2.0)Is it ok to drop derivative information on keywords silently? Should we error (or maybe warn?) if they are nonzero?