Skip to content
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

move Distances, AbstractFFTs and SpecialFunctions' rules out of Zygote #835

Open
2 of 6 tasks
CarloLucibello opened this issue Nov 22, 2020 · 11 comments
Open
2 of 6 tasks
Labels
ChainRules adjoint -> rrule, and further integration

Comments

@CarloLucibello
Copy link
Member

CarloLucibello commented Nov 22, 2020

...to the respective repos, using ChainRulesCore.

This is a list of non-Base functions currently having adjoints in Zygote:

  • Distances
  • StatsFuns
  • SpecialFunctions
  • AbstractFFTs
  • NNlib
  • LoopVectorization
@CarloLucibello
Copy link
Member Author

@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)?

@CarloLucibello CarloLucibello added this to the v0.6 milestone Dec 9, 2020
@DhairyaLGandhi
Copy link
Member

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

@DhairyaLGandhi DhairyaLGandhi removed this from the v0.6 milestone Dec 9, 2020
@willtebbutt
Copy link
Member

willtebbutt commented Dec 9, 2020

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.

@CarloLucibello
Copy link
Member Author

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

@mcabbott
Copy link
Member

mcabbott commented Dec 9, 2020

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.

@willtebbutt
Copy link
Member

willtebbutt commented Dec 9, 2020

@oxinabox what are your thoughts on the AbstractFFTs idea above?

edit: we should consider the solution discussed here

@oxinabox
Copy link
Member

oxinabox commented Dec 9, 2020

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. JuliaMath/AbstractFFTs.jl#47)


For Distances.jl and StatsFuns.jl we should follow JuliaDiff/ChainRules.jl#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.

@mcabbott
Copy link
Member

mcabbott commented Dec 9, 2020

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.

@stevengj
Copy link

stevengj commented Dec 9, 2020

I think it's totally fine for AbstractFFTs to depend on ChainRulesCore…

@devmotion
Copy link
Collaborator

FYI I opened a PR to StatsFuns a while ago (JuliaStats/StatsFuns.jl#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?

@devmotion
Copy link
Collaborator

I made a PR for AbstractFFTs: JuliaMath/AbstractFFTs.jl#58

@mcabbott mcabbott added the ChainRules adjoint -> rrule, and further integration label Jul 22, 2022
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
ChainRules adjoint -> rrule, and further integration
Projects
None yet
Development

No branches or pull requests

7 participants