FluxML / Flux.jl

Relax! Flux is the ML library that doesn't make you tensor
https://fluxml.ai/
Other
4.49k stars 606 forks source link

Zygote fails to differentiate through Flux.params on julia v0.11 #2497

Closed CarloLucibello closed 1 week ago

CarloLucibello commented 1 week ago

With Zygote v0.6.71 on julia v0.11 I get the following error:

julia> m = Chain(Dense(3 => 3), sum);

julia> gradient(m -> sum(Flux.params(m)[1]), m)
ERROR: Can't differentiate foreigncall expression $(Expr(:foreigncall, :(:jl_idset_put_idx), Any, svec(Any, Any, Int64), 0, :(:ccall), %83, %84, %85, %82)).
You might want to check the Zygote limitations documentation.
https://fluxml.ai/Zygote.jl/latest/limitations

Stacktrace:
  [1] error(s::String)
    @ Base ./error.jl:35
  [2] push!
    @ ./idset.jl:62 [inlined]
  [3] (::Zygote.Pullback{Tuple{typeof(push!), IdSet{Any}, typeof(sum)}, Any})(Δ::Nothing)
    @ Zygote ~/.julia/packages/Zygote/Tt5Gx/src/compiler/interface2.jl:0
  [4] params!
    @ ~/.julia/dev/Flux/src/functor.jl:85 [inlined]
  [5] (::Zygote.Pullback{Tuple{typeof(Flux.params!), Params{Zygote.Buffer{…}}, typeof(sum), IdSet{Any}}, Any})(Δ::Nothing)
    @ Zygote ~/.julia/packages/Zygote/Tt5Gx/src/compiler/interface2.jl:0
  [6] params!
    @ ~/.julia/dev/Flux/src/functor.jl:87 [inlined]
  [7] (::Zygote.Pullback{Tuple{typeof(Flux.params!), Params{…}, Tuple{…}, IdSet{…}}, Any})(Δ::Nothing)
    @ Zygote ~/.julia/packages/Zygote/Tt5Gx/src/compiler/interface2.jl:0
  [8] params!
    @ ~/.julia/dev/Flux/src/functor.jl:87 [inlined]
  [9] (::Zygote.Pullback{Tuple{typeof(Flux.params!), Params{Zygote.Buffer{…}}, Chain{Tuple{…}}, IdSet{Any}}, Any})(Δ::Nothing)
    @ Zygote ~/.julia/packages/Zygote/Tt5Gx/src/compiler/interface2.jl:0
 [10] params!
    @ ~/.julia/dev/Flux/src/functor.jl:87 [inlined]
 [11] (::Zygote.Pullback{Tuple{typeof(Flux.params!), Params{Zygote.Buffer{…}}, Tuple{Chain{…}}, IdSet{Any}}, Any})(Δ::Nothing)
    @ Zygote ~/.julia/packages/Zygote/Tt5Gx/src/compiler/interface2.jl:0
 [12] params!
    @ ~/.julia/dev/Flux/src/functor.jl:79 [inlined]
 [13] (::Zygote.Pullback{Tuple{…}, Tuple{…}})(Δ::Nothing)
    @ Zygote ~/.julia/packages/Zygote/Tt5Gx/src/compiler/interface2.jl:0
 [14] params
    @ ~/.julia/dev/Flux/src/functor.jl:127 [inlined]
 [15] (::Zygote.Pullback{Tuple{typeof(Flux.params), Chain{…}}, Tuple{Zygote.Pullback{…}, Zygote.Pullback{…}}})(Δ::Nothing)
    @ Zygote ~/.julia/packages/Zygote/Tt5Gx/src/compiler/interface2.jl:0
 [16] #25
    @ ./REPL[19]:1 [inlined]
 [17] (::Zygote.Pullback{Tuple{…}, Tuple{…}})(Δ::Float32)
    @ Zygote ~/.julia/packages/Zygote/Tt5Gx/src/compiler/interface2.jl:0
 [18] (::Zygote.var"#78#79"{Zygote.Pullback{Tuple{…}, Tuple{…}}})(Δ::Float32)
    @ Zygote ~/.julia/packages/Zygote/Tt5Gx/src/compiler/interface.jl:91
 [19] gradient(f::Function, args::Chain{Tuple{Dense{typeof(identity), Matrix{Float32}, Vector{Float32}}, typeof(sum)}})
    @ Zygote ~/.julia/packages/Zygote/Tt5Gx/src/compiler/interface.jl:148
 [20] top-level scope
    @ REPL[19]:1
Some type information was truncated. Use `show(err)` to see complete types.

Found in #2492

ToucheSir commented 1 week ago

That should be fixed by https://github.com/FluxML/Zygote.jl/pull/1494.

CarloLucibello commented 1 week ago

now fixed in Zygote