search

FluxML / Zygote.jl

21st century AD
https://fluxml.ai/Zygote.jl/
Other
1.48k stars 213 forks source link

Moving adjoints for Distances.jl to the latter as a package extension #1365

Open simsurace opened 1 year ago

simsurace commented 1 year ago

Motivation and description

When trying to precompile function calls to gradients of pairwise, a mutation error is sometimes thrown. According to this discourse thread, this is because the adjoints that are defined to avoid the mutating code are in an @init block, which is not loaded during precompilation.

In the same thread, it has been suggested to use package extensions in order to avoid this problem.

Possible Implementation

No response

devmotion commented 1 year ago

IMO ideally these rules would not be a part of Zygote at all (not even using weak dependencies) but in Distances, using a weak dependency on ChainRulesCore.

devmotion commented 1 year ago

Ah, it seems that this is the intention of the issue? It was not clear to me from the comment above but the title seems to indicate that? In that case, the issue should rather be moved to Distances.jl, shouldn't it? Or just closed in favour of https://github.com/JuliaStats/Distances.jl/issues/172?

simsurace commented 1 year ago

I think that was the suggestion by the participants of the linked discourse thread, but I think the maintainers of Zygote should decide. In case it is moved to Distances.jl, this issue could remain open until the adjoints are removed from Zygote (after the Distances.jl PR is merged).

devmotion commented 1 year ago

Adding things to Distances requires the approval of thr maintainers of Distances, not of Zygote. So to me this seems to be mainly a Distances issue? Only removing rules from Zygote if they are available in Distances would require the approval of Zygote maintainers.

CarloLucibello commented 1 year ago

On the Zygote side, it would be nice to avoid a breaking change. Maybe the extension in Distances.jl could depend on Zygote being loaded? This is not ideal though. I wonder if an extension depending on package A is also activated when using B with B loading A. In that case depending on ChainRulesCore would be fine.

devmotion commented 1 year ago

Maybe the extension in Distances.jl could depend on Zygote being loaded?

Hopefully not. ChainRulesCore is a tiny (weak) dependency whereas Zygote would be huge.

I wonder if an extension depending on package A is also activated when using B with B loading A.

As far as I know that's the case. It is like a (much) better Requires that also supports precompilation and extensions depending on multiple weak dependencies.

CarloLucibello commented 1 year ago

As far as I know that's the case.

true, confirmed in https://github.com/JuliaLang/Pkg.jl/issues/3348

ToucheSir commented 1 year ago

Unless Distances.jl is fine with adding Requires for pre-1.9 support, I think we're stuck with maintaining a copy of these rules. It shouldn't be the end of the world (i.e. whenever the package extension changes, we should be able to run a couple sed commands on the source and copy it here), but if it becomes a burden we can always tell users who need the rules to switch to 1.9+ instead of updating the older copy here.

devmotion commented 1 year ago

FYI I opened a PR to get feedback from maintainers of Distances: https://github.com/JuliaStats/Distances.jl/pull/246