Zygote.jl
Zygote.jl copied to clipboard
move Distances, AbstractFFTs and SpecialFunctions' rules out of Zygote
...to the respective repos, using ChainRulesCore.
This is a list of non-Base functions currently having adjoints in Zygote:
- [ ] Distances
- [ ] StatsFuns
- [ ] SpecialFunctions
- [ ] AbstractFFTs
- [x] NNlib
- [x] LoopVectorization
@mcabbott @willtebbutt so we want to do all this by v0.6, migrating stuff to packages and ChainRules if package owners are not willing to carry the extra dependence (which could be the case for all of them, since they are all very lean)?
Wouldn't put it out for 0.6, but bump it for a bit later since this is a fair amount of work and can incur breakages which need to be accounted for beforehand
My two cent FWIW: I wouldn't personally object to these all going into 0.6, provided that new patch releases to 0.5 aren't blocked while this is happening (i.e. the development happens on a separate 0.6-DEV branch). I guess it depends on what time scale you want to see 0.6 on.
Also, I think SpecialFunctions is being handled already. See this PR and the corresponding change in the ChainRules code (which actually might be a way to push making a breaking release down the line @mcabbott @simeonschaub?) Do you know what SpecialFunctions-related rules we have in Zygote?
edit: looks like SpecialFunctions only appears in forward through DiffRules, so shoudn't be too hard to remove.
NNlib and LoopVectorization excision are already on master, so we could just do whatever is left to do for SpecialFunctions and tag v0.6, without the additional complexity of having to branch out a release-0.5 and start backporting things
Does Zygote's forward mode use ChainRules? If so, then removing SpecialFunctions should be easy.
Which others might be worth trying to do quickly, and holding up 0.6?
If it's just a short time we should just leave everything on master; a medium time maybe it's worth messing with branches; a long time then it should just be 0.7 someday, I think. Messing with Requires to make temporary solutions doesn't sound like great use of everyone's time.
-
Distances.jl sounds unlikely in the short term. Currently loaded by Requires. One idea is to move this Requires dep to ChainRules -- if not ideal, still a step forward. That can be done without breaking.
-
AbstractFFTs is a very light stub package, much smaller than ChainRulesCore I think. So you could argue that ChainRules should just depend on that. In which case the move from status quo can be done without breaking.
-
That leaves StatsFuns.jl. I don't see an issue, but perhaps that one is worth trying? If its owners are opposed, then it would be exactly like Distances.jl -- also currently loaded by Requires.
@oxinabox what are your thoughts on the AbstractFFTs idea above?
edit: we should consider the solution discussed here
Does Zygote's forward mode use ChainRules? If so, then removing SpecialFunctions should be easy.
Sadly it does not.
Re: AbstractFFTs. We should ask @StevenGj what to do. If ever there was a package that is is ok for ChainRules to directly depend on it is probably that one. It's almost a standard library and has only had one breaking change in the last 3 years (Not sure why it hasn't tagged 1.0. https://github.com/JuliaMath/AbstractFFTs.jl/issues/47)
For Distances.jl and StatsFuns.jl we should follow https://github.com/JuliaDiff/ChainRules.jl/issues/280 which properly lets the resolver handle things. Unlike just using Requires alone. With the cost and the benifit, that things will actually be held back. Not letting the resolver to its work is asking for trouble in the long term. Maybe we should do that for AbstractFFT too.
OK. It sounds like this ChainRules_Distances_Glue.jl idea would work, if someone is happy to set it up. (Maybe ChainRules is a good candidate for that mult-package-one-repository story?)
The transition to that can be done without breaking, so 0.6 need not wait for it.
I think it's totally fine for AbstractFFTs to depend on ChainRulesCore…
FYI I opened a PR to StatsFuns a while ago (https://github.com/JuliaStats/StatsFuns.jl/pull/106) but it has not been approved (yet) even though the package already depends on ChainRulesCore indirectly via SpecialFunctions. Maybe one should create a glue package instead and add the definitions there?
I made a PR for AbstractFFTs: https://github.com/JuliaMath/AbstractFFTs.jl/pull/58