Zygote is broken on Julia 1.10

[mcabbott]

On master, Zygote cannot be loaded. Commenting out precompile, what works and doesn't looks like:

julia> gradient(log, 1.0)
(1.0,)

julia> gradient(x -> log(x), 1.0)
ERROR: MethodError: no method matching iterate(::Nothing)
Stacktrace:
 [1] indexed_iterate(I::Nothing, i::Int64)
   @ Base ./tuple.jl:93
 [2] chain_rrule
   @ ~/.julia/dev/Zygote/src/compiler/chainrules.jl:223 [inlined]
 [3] macro expansion
   @ ~/.julia/dev/Zygote/src/compiler/interface2.jl:0 [inlined]
 [4] _pullback(ctx::Zygote.Context{false}, f::var"#15#16", args::Float64)
   @ Zygote ~/.julia/dev/Zygote/src/compiler/interface2.jl:9
 [5] pullback(f::Function, cx::Zygote.Context{false}, args::Float64)
   @ Zygote ~/.julia/dev/Zygote/src/compiler/interface.jl:44
 [6] pullback
   @ ~/.julia/dev/Zygote/src/compiler/interface.jl:42 [inlined]
 [7] gradient(f::Function, args::Float64)
   @ Zygote ~/.julia/dev/Zygote/src/compiler/interface.jl:96

julia> gradient(sum, [1,2])
([1.0, 1.0],)

julia> gradient(abs2∘sum, [1,2])
ERROR: MethodError: no method matching iterate(::Nothing)

julia> versioninfo()
Julia Version 1.10.0-DEV.918
Commit fe1e1c4bb4 (2023-03-30 18:36 UTC)
Platform Info:
  OS: macOS (arm64-apple-darwin21.6.0)
  CPU: 8 × Apple M1

[ToucheSir]

Zygote.jl/src/compiler/chainrules.jl

Line 223 in 413728b

y, back = rrule(config, f, args...)

fails trying to destructure a nothing return from rrule. That in turn implies

Zygote.jl/src/compiler/interface2.jl

Lines 20 to 21 in 413728b

	hascr, cr_edge = has_chain_rrule(cr_T)
	hascr && return :($chain_rrule_f(ZygoteRuleConfig(ctx), f, args...))

has erroneously succeeded. My guess is that we're hitting

Zygote.jl/src/compiler/chainrules.jl

Line 73 in 413728b

return true, nothing

. I don't think we can blame ambiguities here, so something about the earlier logic is off. cc @oxinabox for possible ideas.

[oxinabox]

just fyi, I am aware of this issue, it's in my backlog.
But not in any hurry to get to it since 1.9 isn't out so we are a long way from 1.10.

[maleadt]

we are a long way from 1.10.

It would also help for PkgEval; quite some packages are failing (and getting blacklisted) because of this.

[IanButterworth]

Does anyone have a workaround for this? I'm trying to use the new profiling on master but this is a blocker

[IanButterworth]

Adding some shows master...IanButterworth:Zygote.jl:ib/debug
1.8.5

julia> gradient(x->log(x), 1.0)
(config_T, arg_Ts) = Iterators.peel(T.parameters) = (Zygote.ZygoteRuleConfig{Zygote.Context{false}}, Base.Iterators.Rest{Core.SimpleVector, Int64}(svec(Zygote.ZygoteRuleConfig{Zygote.Context{false}}, var"#3#4", Float64), 2))
configured_rrule_m = meta(Tuple{typeof(rrule), config_T, arg_Ts...}) = Metadata for rrule(::ChainRulesCore.RuleConfig, args...) in ChainRulesCore at /home/ian/.julia/packages/ChainRulesCore/a4mIA/src/rules.jl:134
is_ambig = false
_is_rrule_redispatcher(configured_rrule_m.method) = true
no_rrule_m = Metadata for no_rrule(::Any, ...) in ChainRulesCore at /home/ian/.julia/packages/ChainRulesCore/a4mIA/src/rules.jl:190
rrule_m = Metadata for rrule(::Any, ...) in ChainRulesCore at /home/ian/.julia/packages/ChainRulesCore/a4mIA/src/rules.jl:131
(config_T, arg_Ts) = Iterators.peel(T.parameters) = (Zygote.ZygoteRuleConfig{Zygote.Context{false}}, Base.Iterators.Rest{Core.SimpleVector, Int64}(svec(Zygote.ZygoteRuleConfig{Zygote.Context{false}}, typeof(log), Float64), 2))
configured_rrule_m = meta(Tuple{typeof(rrule), config_T, arg_Ts...}) = Metadata for rrule(::ChainRulesCore.RuleConfig, args...) in ChainRulesCore at /home/ian/.julia/packages/ChainRulesCore/a4mIA/src/rules.jl:134
is_ambig = false
_is_rrule_redispatcher(configured_rrule_m.method) = true
no_rrule_m = Metadata for no_rrule(::Any, ...) in ChainRulesCore at /home/ian/.julia/packages/ChainRulesCore/a4mIA/src/rules.jl:190
rrule_m = Metadata for rrule(::typeof(log), x::Number) in ChainRules at /home/ian/.julia/packages/ChainRules/bEtjZ/src/rulesets/Base/fastmath_able.jl:60
(1.0,)

1.10

julia> gradient(x->log(x), 1.0)
(config_T, arg_Ts) = Iterators.peel(T.parameters) = (Zygote.ZygoteRuleConfig{Zygote.Context{false}}, Base.Iterators.Rest{Core.SimpleVector, Int64}(svec(Zygote.ZygoteRuleConfig{Zygote.Context{false}}, var"#3#4", Float64), 2))
configured_rrule_m = meta(Tuple{typeof(rrule), config_T, arg_Ts...}) = nothing
is_ambig = true
(config_T, arg_Ts) = Iterators.peel(T.parameters) = (Zygote.ZygoteRuleConfig{Zygote.Context{false}}, Base.Iterators.Rest{Core.SimpleVector, Int64}(svec(Zygote.ZygoteRuleConfig{Zygote.Context{false}}, var"#3#4", Float64), 2))
configured_rrule_m = meta(Tuple{typeof(rrule), config_T, arg_Ts...}) = nothing
is_ambig = true
ERROR: type Nothing has no field method
...

So the first deviation is at meta(Tuple{typeof(rrule), config_T, arg_Ts...}) which returns nothing on 1.10

[ToucheSir]

Thanks for investigating, Ian. This kind of failure has happened before, and unfortunately will require digging into the IRTools side to get a proper idea of what went wrong. I wonder if we shouldn't vendor this particular bit of code into Zygote to make future maintenance easier and possibly improve error handling (i.e. more informative errors and fewer silent fallbacks to nothing).

[IanButterworth]

I don't think this is the issue because there's no error in IRTools, but it smells like it's adjacent FluxML/IRTools.jl#109

[ToucheSir]

IRTools may be swallowing the error or some other invalid state, but I haven't looked into that.

[IanButterworth]

IRTools isn't very happy at all on 1.10 FluxML/IRTools.jl#111

[maleadt]

The problem is world age related. The @generated _pullback discards the codegen world, which needs to be manually propagated after JuliaLang/julia#48766. Sadly, this is relatively invasive, and requires manual expansion of the @generated macro. Then, in combination with FluxML/IRTools.jl#111, the correct world needs to be passed to IRTools.meta.

Initial fixes in https://github.com/maleadt/Zygote.jl/tree/dev, depending on FluxML/IRTools.jl#111.