Add weak dependency on ChainRulesCore #246
Conversation
Codecov ReportAll modified lines are covered by tests ✅
... and 1 file with indirect coverage changes 📢 Thoughts on this report? Let us know!. |
|
With extensions I think it is fine to add code like this even to light weight packages. I do think that JuliaStats/LogExpFunctions.jl#64 should be taken into account. |
|
Is there anything blocking this? |
|
@KristofferC, could this be merged? |
|
Bump 🙂 |
devmotion
changed the title
|
|
||
| import ChainRulesCore | ||
|
|
||
| const CRC = ChainRulesCore |
There was a problem hiding this comment.
Why not do import ChainRulesCore: rrule and using ChainRulesCore: RuleConfig, ... instead?
There was a problem hiding this comment.
I'm not sure, it was a while ago. I generally prefer explicit namespace declarations over import statements, but I assume I didn't want to write ChainRulesCore all the time and therefore added the alias (which I have seen in other packages and issues as well).
I wonder if maybe @sethaxen or @willtebbutt (who contributed to the rules in Zygote multiple times) would have time to check the rules? |
There was a problem hiding this comment.
I think this broadly looks good.
Would you consider incorporating tests similar to those found here: https://github.com/FluxML/Zygote.jl/blob/29fa32a688fb4454d2ee81b9ba2a6484e4468bda/test/gradcheck.jl#L1212 ?
IIRC they came up in the context of some GP stuff, so it would be great to ensure that we don't regress.
|
Bump! |
|
I'll return to the PR after my vacation, the PR is high up on my todo-list 🙂 |
|
I added tests with repeated columns: 2e8563a Is this what you had in mind @willtebbutt? |
|
Bump 🙂 |
|
Good to go, @nalimilan? |
This PR proposes to add a weak dependency on ChainRulesCore on Julia versions with support for extensions (i.e., Julia >= 1.9).
This would fix #172 (and FluxML/Zygote.jl#1365) without affecting load and compilation times for users that do not load ChainRulesCore.
The initial set of reverse-mode rules in this PR covers the existing rules in Zygote and hence could replace these in Julia >= 1.9: Zygote could, e.g., wrap the Requires-block for Distances in a
if !isdefined(Base, :get_extension)conditional.The implementation in this PR is different from the code in FluxML/Zygote.jl#923 since tests fail with the rules in the Zygote PR (e.g., due to type inference issues).
The rules in this PR are also all explicit and do not call back into the AD system, in contrast to the current implementation in Zygote.
IMO there are two main motivations for moving the rules to Distances.jl: It allows other ChainRules-compatible AD systems to use them as well without having to depend on Zygote and it fixes precompilation issues reported in FluxML/Zygote.jl#1365. The second point could be fixed also by making Distances a weak dependency of Zygote but that would not solve the first problem.
I know that maintainers of Distances have tried to keep the package lightweight and to not add any dependencies that would increase load times. However, I think the support for weak dependencies in Pkg addresses these concerns. In a few other packages ChainRules-definitions were already moved to extensions recently, e.g., JuliaMath/ChangesOfVariables.jl#12 and JuliaStats/LogExpFunctions.jl#62.