Conversation
src/zygoterules.jl
Outdated
| # (:σ, :(conj.(Ω .* (1 .- Ω)))), | ||
| # ] | ||
| # pullback = Symbol(:broadcasted_, f, :_pullback) | ||
| # @eval @adjoint function Base.Broadcast.broadcasted(::typeof($f), x::Numeric) |
There was a problem hiding this comment.
I am not sure this is a good idea, since IIUC, loading NNlib will now always error if Zygote is not loaded first. Perhaps we can use Requires.jl for this instead?
There was a problem hiding this comment.
I don't understand, when should it error? I can do using NNlib just fine
There was a problem hiding this comment.
Using requires is not an option I think, because Zygote v0.5 precompilation triggers NNlib's one, NNlib define the rules, then Zygote tries to define the same rules and gives the warning
There was a problem hiding this comment.
I don't understand, when should it error? I can do
using NNlibjust fine
I can't try it ATM, but since it doesn't load ZygoteRules now, I don't know where it should get the definition of @adjoint from. It might just work for the tests, since those already load Zygote. Did you do using NNlib before loading any other packages?
There was a problem hiding this comment.
Using requires is not an option I think, because Zygote v0.5 precompilation triggers NNlib's one, NNlib define the rules, then Zygote tries to define the same rules and gives the warning
I think we could use requires for checking if Zygote is loaded and then only define those rules if it's the appropriate version.
There was a problem hiding this comment.
I don't know where it should get the definition of @adjoint from
ah, maybe you didn't see the new rules down below?
I think we could use requires for checking if Zygote is loaded and then only define those rules if it's the appropriate version.
I think this won't work because Zygote v0.5 depends on NNlib, so NNlib is always loaded first
There was a problem hiding this comment.
I don't know where it should get the definition of
@adjointfrom
maybe there is been a misunderstanding, I replaced ZygoteRules' adjoint with ChainRules' rrules. I don't know why it wasn't done before, maybe because is not something any AD engine would want to do
There was a problem hiding this comment.
🤦I didn't see this was commented out for some reason, sorry! If this is just a temporary workaround, using ChainRulesCore for these seems ok to me, but we really want to avoid this once people are actually using other ADs than Zygote. JuliaDiff/ChainRulesCore.jl#270 might be helpful here.
|
@oxinabox Will this cause problems for Nabla.jl? |
|
No, this should be great with Nabla. |
Co-authored-by: Simeon Schaub <simeondavidschaub99@gmail.com>
3 participants
With this change, Flux's tests pass with both Zygote v0.5 and v0.6,
redundancy or ChainRule's rrules here with the ones defined in Zygote v0.5 doesn't seem to be a problem.
ZygoteRules' adjoints instead were giving redefinition warnings.
With this change we can tag a v0.7.10 release instead of v0.8, which would relieve us of a lot of pain:
CUDA depends on NNlib, and Flux depends on CUDA,
but CUDA master is only compatible with julia 1.6, so it would take a while to see CUDA release with the NNlib's compat bound update. Moreover, we would have to drop julia < 1.6 support in Flux.
@simeonschaub do you think this change is fine?