Zygote borks on simple Turing model

[dlakelan]

I'm building a Sysimage using PackageCompiler to do data analysis without the TTFP issues.

When I include the following in the precompile exerciser script, it borks:

@model function foo(data)
    n ~ Gamma(3,2)
    foo ~ arraydist([Gamma(3,2) for i in 1:length(data)])
    bar ~ MvNormal([0.0 for i in 1:length(data)],1.0)
    baz ~ Normal(3.0,2.0)
end

Turing.setadbackend(:reversediff)
Turing.setrdcache(true)
Turing.emptyrdcache()



ch = sample(foo(collect(1:100)),NUTS(1000,.8),2000)
summarize(ch)

Turing.setadbackend(:zygote)
ch = sample(foo(collect(1:100)),NUTS(1000,.8),2000)
summary(ch)

The first sample using the reversediff backend works fine, but it crashes out on the second sample call using the Zygote backend with a backtrace like this:

ERROR: LoadError: MethodError: no method matching +(::NamedTuple{(:x,), Tuple{Float64}}, ::ChainRulesCore.Tangent{Base.RefValue{Float64}, NamedTuple{(:x,), Tuple{Float64}}})
Closest candidates are:
  +(::Any, ::Any, ::Any, ::Any...) at operators.jl:655
  +(::Union{InitialValues.NonspecificInitialValue, InitialValues.SpecificInitialValue{typeof(+)}}, ::Any) at /home/dlakelan/.julia/packages/InitialValues/P5PLf/src/InitialValues.jl:153
  +(::ChainRulesCore.AbstractThunk, ::ChainRulesCore.Tangent) at /home/dlakelan/.julia/packages/ChainRulesCore/Y1Mee/src/tangent_arithmetic.jl:121
  ...
Stacktrace:
  [1] accum(x::NamedTuple{(:x,), Tuple{Float64}}, y::ChainRulesCore.Tangent{Base.RefValue{Float64}, NamedTuple{(:x,), Tuple{Float64}}})
    @ Zygote ~/.julia/packages/Zygote/rv6db/src/lib/lib.jl:17
  [2] _broadcast_getindex_evalf
    @ ./broadcast.jl:670 [inlined]
  [3] _broadcast_getindex
    @ ./broadcast.jl:643 [inlined]
  [4] getindex
    @ ./broadcast.jl:597 [inlined]
  [5] copy
    @ ./broadcast.jl:943 [inlined]
  [6] materialize
    @ ./broadcast.jl:904 [inlined]
  [7] accum(x::Vector{NamedTuple{(:x,), Tuple{Float64}}}, ys::Vector{ChainRulesCore.AbstractTangent})
    @ Zygote ~/.julia/packages/Zygote/rv6db/src/lib/lib.jl:25
  [8] macro expansion
    @ ~/.julia/packages/Zygote/rv6db/src/lib/lib.jl:27 [inlined]
  [9] accum(x::NamedTuple{(:varinfo, :logps), Tuple{Nothing, Vector{NamedTuple{(:x,), Tuple{Float64}}}}}, y::NamedTuple{(:varinfo, :logps), Tuple{NamedTuple{(:metadata, :logp, :num_produce), Tuple{NamedTuple{(:n, :foo, :bar, :baz), NTuple{4, NamedTuple{(:idcs, :vns, :ranges, :vals, :dists, :gids, :orders, :flags), Tuple{Nothing, Nothing, Nothing, Vector{Float64}, Nothing, Nothing, Nothing, Nothing}}}}, Nothing, Nothing}}, Vector{ChainRulesCore.AbstractTangent}}})
    @ Zygote ~/.julia/packages/Zygote/rv6db/src/lib/lib.jl:27
 [10] getindex
    @ ./tuple.jl:29 [inlined]
 [11] gradindex
    @ ~/.julia/packages/Zygote/rv6db/src/compiler/reverse.jl:12 [inlined]
 [12] Pullback
    @ ~/.julia/packages/DynamicPPL/RcfQU/src/model.jl:438 [inlined]
 [13] (::typeof(∂(evaluate_threadsafe)))(Δ::Nothing)
    @ Zygote ~/.julia/packages/Zygote/rv6db/src/compiler/interface2.jl:0
 [14] Pullback
    @ ~/.julia/packages/DynamicPPL/RcfQU/src/model.jl:391 [inlined]
 [15] (::typeof(∂(λ)))(Δ::Nothing)
    @ Zygote ~/.julia/packages/Zygote/rv6db/src/compiler/interface2.jl:0
 [16] Pullback
    @ ~/.julia/packages/DynamicPPL/RcfQU/src/model.jl:383 [inlined]
 [17] (::typeof(∂(λ)))(Δ::Nothing)
    @ Zygote ~/.julia/packages/Zygote/rv6db/src/compiler/interface2.jl:0
 [18] #209
    @ ~/.julia/packages/Zygote/rv6db/src/lib/lib.jl:203 [inlined]
 [19] #1740#back
    @ ~/.julia/packages/ZygoteRules/AIbCs/src/adjoint.jl:67 [inlined]
 [20] Pullback
    @ ~/.julia/packages/DynamicPPL/RcfQU/src/model.jl:396 [inlined]
 [21] (::typeof(∂(λ)))(Δ::Nothing)
    @ Zygote ~/.julia/packages/Zygote/rv6db/src/compiler/interface2.jl:0
 [22] Pullback
    @ ~/.julia/packages/Turing/uMQmD/src/core/ad.jl:165 [inlined]
 [23] (::typeof(∂(λ)))(Δ::Int64)
    @ Zygote ~/.julia/packages/Zygote/rv6db/src/compiler/interface2.jl:0
 [24] (::Zygote.var"#52#53"{typeof(∂(λ))})(Δ::Int64)
    @ Zygote ~/.julia/packages/Zygote/rv6db/src/compiler/interface.jl:41
 [25] gradient_logp(backend::Turing.Core.ZygoteAD, θ::Vector{Float64}, vi::DynamicPPL.TypedVarInfo{NamedTuple{(:n, :foo, :bar, :baz), Tuple{DynamicPPL.Metadata{Dict{AbstractPPL.VarName{:n, Tuple{}}, Int64}, Vector{Gamma{Float64}}, Vector{AbstractPPL.VarName{:n, Tuple{}}}, Vector{Float64}, Vector{Set{DynamicPPL.Selector}}}, DynamicPPL.Metadata{Dict{AbstractPPL.VarName{:foo, Tuple{}}, Int64}, Vector{Product{Continuous, Gamma{Float64}, Vector{Gamma{Float64}}}}, Vector{AbstractPPL.VarName{:foo, Tuple{}}}, Vector{Float64}, Vector{Set{DynamicPPL.Selector}}}, DynamicPPL.Metadata{Dict{AbstractPPL.VarName{:bar, Tuple{}}, Int64}, Vector{IsoNormal}, Vector{AbstractPPL.VarName{:bar, Tuple{}}}, Vector{Float64}, Vector{Set{DynamicPPL.Selector}}}, DynamicPPL.Metadata{Dict{AbstractPPL.VarName{:baz, Tuple{}}, Int64}, Vector{Normal{Float64}}, Vector{AbstractPPL.VarName{:baz, Tuple{}}}, Vector{Float64}, Vector{Set{DynamicPPL.Selector}}}}}, Float64}, model::DynamicPPL.Model{typeof(foo), (:data,), (), (), Tuple{Vector{Int64}}, Tuple{}, DynamicPPL.DefaultContext}, sampler::DynamicPPL.Sampler{NUTS{Turing.Core.ZygoteAD, (), AdvancedHMC.DiagEuclideanMetric}}, context::DynamicPPL.DefaultContext)
    @ Turing.Core ~/.julia/packages/Turing/uMQmD/src/core/ad.jl:171
 [26] gradient_logp (repeats 2 times)
    @ ~/.julia/packages/Turing/uMQmD/src/core/ad.jl:83 [inlined]
 [27] ∂logπ∂θ
    @ ~/.julia/packages/Turing/uMQmD/src/inference/hmc.jl:433 [inlined]
 [28] ∂H∂θ
    @ ~/.julia/packages/AdvancedHMC/HQHnm/src/hamiltonian.jl:31 [inlined]
 [29] phasepoint
    @ ~/.julia/packages/AdvancedHMC/HQHnm/src/hamiltonian.jl:76 [inlined]
 [30] phasepoint(rng::Random._GLOBAL_RNG, θ::Vector{Float64}, h::AdvancedHMC.Hamiltonian{AdvancedHMC.DiagEuclideanMetric{Float64, Vector{Float64}}, Turing.Inference.var"#logπ#54"{DynamicPPL.TypedVarInfo{NamedTuple{(:n, :foo, :bar, :baz), Tuple{DynamicPPL.Metadata{Dict{AbstractPPL.VarName{:n, Tuple{}}, Int64}, Vector{Gamma{Float64}}, Vector{AbstractPPL.VarName{:n, Tuple{}}}, Vector{Float64}, Vector{Set{DynamicPPL.Selector}}}, DynamicPPL.Metadata{Dict{AbstractPPL.VarName{:foo, Tuple{}}, Int64}, Vector{Product{Continuous, Gamma{Float64}, Vector{Gamma{Float64}}}}, Vector{AbstractPPL.VarName{:foo, Tuple{}}}, Vector{Float64}, Vector{Set{DynamicPPL.Selector}}}, DynamicPPL.Metadata{Dict{AbstractPPL.VarName{:bar, Tuple{}}, Int64}, Vector{IsoNormal}, Vector{AbstractPPL.VarName{:bar, Tuple{}}}, Vector{Float64}, Vector{Set{DynamicPPL.Selector}}}, DynamicPPL.Metadata{Dict{AbstractPPL.VarName{:baz, Tuple{}}, Int64}, Vector{Normal{Float64}}, Vector{AbstractPPL.VarName{:baz, Tuple{}}}, Vector{Float64}, Vector{Set{DynamicPPL.Selector}}}}}, Float64}, DynamicPPL.Sampler{NUTS{Turing.Core.ZygoteAD, (), AdvancedHMC.DiagEuclideanMetric}}, DynamicPPL.Model{typeof(foo), (:data,), (), (), Tuple{Vector{Int64}}, Tuple{}, DynamicPPL.DefaultContext}}, Turing.Inference.var"#∂logπ∂θ#53"{DynamicPPL.TypedVarInfo{NamedTuple{(:n, :foo, :bar, :baz), Tuple{DynamicPPL.Metadata{Dict{AbstractPPL.VarName{:n, Tuple{}}, Int64}, Vector{Gamma{Float64}}, Vector{AbstractPPL.VarName{:n, Tuple{}}}, Vector{Float64}, Vector{Set{DynamicPPL.Selector}}}, DynamicPPL.Metadata{Dict{AbstractPPL.VarName{:foo, Tuple{}}, Int64}, Vector{Product{Continuous, Gamma{Float64}, Vector{Gamma{Float64}}}}, Vector{AbstractPPL.VarName{:foo, Tuple{}}}, Vector{Float64}, Vector{Set{DynamicPPL.Selector}}}, DynamicPPL.Metadata{Dict{AbstractPPL.VarName{:bar, Tuple{}}, Int64}, Vector{IsoNormal}, Vector{AbstractPPL.VarName{:bar, Tuple{}}}, Vector{Float64}, Vector{Set{DynamicPPL.Selector}}}, DynamicPPL.Metadata{Dict{AbstractPPL.VarName{:baz, Tuple{}}, Int64}, Vector{Normal{Float64}}, Vector{AbstractPPL.VarName{:baz, Tuple{}}}, Vector{Float64}, Vector{Set{DynamicPPL.Selector}}}}}, Float64}, DynamicPPL.Sampler{NUTS{Turing.Core.ZygoteAD, (), AdvancedHMC.DiagEuclideanMetric}}, DynamicPPL.Model{typeof(foo), (:data,), (), (), Tuple{Vector{Int64}}, Tuple{}, DynamicPPL.DefaultContext}}})
    @ AdvancedHMC ~/.julia/packages/AdvancedHMC/HQHnm/src/hamiltonian.jl:153
 [31] initialstep(rng::Random._GLOBAL_RNG, model::DynamicPPL.Model{typeof(foo), (:data,), (), (), Tuple{Vector{Int64}}, Tuple{}, DynamicPPL.DefaultContext}, spl::DynamicPPL.Sampler{NUTS{Turing.Core.ZygoteAD, (), AdvancedHMC.DiagEuclideanMetric}}, vi::DynamicPPL.TypedVarInfo{NamedTuple{(:n, :foo, :bar, :baz), Tuple{DynamicPPL.Metadata{Dict{AbstractPPL.VarName{:n, Tuple{}}, Int64}, Vector{Gamma{Float64}}, Vector{AbstractPPL.VarName{:n, Tuple{}}}, Vector{Float64}, Vector{Set{DynamicPPL.Selector}}}, DynamicPPL.Metadata{Dict{AbstractPPL.VarName{:foo, Tuple{}}, Int64}, Vector{Product{Continuous, Gamma{Float64}, Vector{Gamma{Float64}}}}, Vector{AbstractPPL.VarName{:foo, Tuple{}}}, Vector{Float64}, Vector{Set{DynamicPPL.Selector}}}, DynamicPPL.Metadata{Dict{AbstractPPL.VarName{:bar, Tuple{}}, Int64}, Vector{IsoNormal}, Vector{AbstractPPL.VarName{:bar, Tuple{}}}, Vector{Float64}, Vector{Set{DynamicPPL.Selector}}}, DynamicPPL.Metadata{Dict{AbstractPPL.VarName{:baz, Tuple{}}, Int64}, Vector{Normal{Float64}}, Vector{AbstractPPL.VarName{:baz, Tuple{}}}, Vector{Float64}, Vector{Set{DynamicPPL.Selector}}}}}, Float64}; init_params::Nothing, nadapts::Int64, kwargs::Base.Pairs{Symbol, Union{}, Tuple{}, NamedTuple{(), Tuple{}}})
    @ Turing.Inference ~/.julia/packages/Turing/uMQmD/src/inference/hmc.jl:167
 [32] step(rng::Random._GLOBAL_RNG, model::DynamicPPL.Model{typeof(foo), (:data,), (), (), Tuple{Vector{Int64}}, Tuple{}, DynamicPPL.DefaultContext}, spl::DynamicPPL.Sampler{NUTS{Turing.Core.ZygoteAD, (), AdvancedHMC.DiagEuclideanMetric}}; resume_from::Nothing, kwargs::Base.Pairs{Symbol, Int64, Tuple{Symbol}, NamedTuple{(:nadapts,), Tuple{Int64}}})
    @ DynamicPPL ~/.julia/packages/DynamicPPL/RcfQU/src/sampler.jl:87
 [33] macro expansion
    @ ~/.julia/packages/AbstractMCMC/BPJCW/src/sample.jl:123 [inlined]
 [34] macro expansion
    @ ~/.julia/packages/ProgressLogging/6KXlp/src/ProgressLogging.jl:328 [inlined]
 [35] (::AbstractMCMC.var"#21#22"{Bool, String, Nothing, Int64, Int64, Base.Pairs{Symbol, Int64, Tuple{Symbol}, NamedTuple{(:nadapts,), Tuple{Int64}}}, Random._GLOBAL_RNG, DynamicPPL.Model{typeof(foo), (:data,), (), (), Tuple{Vector{Int64}}, Tuple{}, DynamicPPL.DefaultContext}, DynamicPPL.Sampler{NUTS{Turing.Core.ZygoteAD, (), AdvancedHMC.DiagEuclideanMetric}}, Int64, Int64})()
    @ AbstractMCMC ~/.julia/packages/AbstractMCMC/BPJCW/src/logging.jl:11
 [36] with_logstate(f::Function, logstate::Any)
    @ Base.CoreLogging ./logging.jl:511
 [37] with_logger(f::Function, logger::LoggingExtras.TeeLogger{Tuple{LoggingExtras.EarlyFilteredLogger{TerminalLoggers.TerminalLogger, AbstractMCMC.var"#1#3"{Module}}, LoggingExtras.EarlyFilteredLogger{Logging.ConsoleLogger, AbstractMCMC.var"#2#4"{Module}}}})
    @ Base.CoreLogging ./logging.jl:623
 [38] with_progresslogger(f::Function, _module::Module, logger::Logging.ConsoleLogger)
    @ AbstractMCMC ~/.julia/packages/AbstractMCMC/BPJCW/src/logging.jl:34
 [39] macro expansion
    @ ~/.julia/packages/AbstractMCMC/BPJCW/src/logging.jl:10 [inlined]
 [40] mcmcsample(rng::Random._GLOBAL_RNG, model::DynamicPPL.Model{typeof(foo), (:data,), (), (), Tuple{Vector{Int64}}, Tuple{}, DynamicPPL.DefaultContext}, sampler::DynamicPPL.Sampler{NUTS{Turing.Core.ZygoteAD, (), AdvancedHMC.DiagEuclideanMetric}}, N::Int64; progress::Bool, progressname::String, callback::Nothing, discard_initial::Int64, thinning::Int64, chain_type::Type, kwargs::Base.Pairs{Symbol, Int64, Tuple{Symbol}, NamedTuple{(:nadapts,), Tuple{Int64}}})
    @ AbstractMCMC ~/.julia/packages/AbstractMCMC/BPJCW/src/sample.jl:114
 [41] sample(rng::Random._GLOBAL_RNG, model::DynamicPPL.Model{typeof(foo), (:data,), (), (), Tuple{Vector{Int64}}, Tuple{}, DynamicPPL.DefaultContext}, sampler::DynamicPPL.Sampler{NUTS{Turing.Core.ZygoteAD, (), AdvancedHMC.DiagEuclideanMetric}}, N::Int64; chain_type::Type, resume_from::Nothing, progress::Bool, nadapts::Int64, discard_adapt::Bool, discard_initial::Int64, kwargs::Base.Pairs{Symbol, Union{}, Tuple{}, NamedTuple{(), Tuple{}}})
    @ Turing.Inference ~/.julia/packages/Turing/uMQmD/src/inference/hmc.jl:133
 [42] sample
    @ ~/.julia/packages/Turing/uMQmD/src/inference/hmc.jl:116 [inlined]
 [43] #sample#2
    @ ~/.julia/packages/Turing/uMQmD/src/inference/Inference.jl:142 [inlined]
 [44] sample
    @ ~/.julia/packages/Turing/uMQmD/src/inference/Inference.jl:142 [inlined]
 [45] #sample#1
    @ ~/.julia/packages/Turing/uMQmD/src/inference/Inference.jl:132 [inlined]
 [46] sample(model::DynamicPPL.Model{typeof(foo), (:data,), (), (), Tuple{Vector{Int64}}, Tuple{}, DynamicPPL.DefaultContext}, alg::NUTS{Turing.Core.ZygoteAD, (), AdvancedHMC.DiagEuclideanMetric}, N::Int64)
    @ Turing.Inference ~/.julia/packages/Turing/uMQmD/src/inference/Inference.jl:132
 [47] top-level scope
    @ ~/juliapkgimage/dataanalys.jl:77
in expression starting at /home/dlakelan/juliapkgimage/dataanalys.jl:77
ERROR: LoadError: failed process: Process(`/home/dlakelan/julia-latest/bin/julia --color=yes --startup-file=no --sysimage=/home/dlakelan/julia-latest/lib/julia/sys.so --project=/home/dlakelan/juliapkgimage --compile=all --trace-compile=/tmp/jl_EI8YlN /home/dlakelan/juliapkgimage/dataanalys.jl`, ProcessExited(1)) [1]

[DhairyaLGandhi]

It seems like there are a lot of leaking types, a Tangent should not be constructed with a RefValue as in the stacktrace. Otherwise, accumulation should support named tuples and tangents. Maybe a better reduced mwe could be accum(::NamedTuple, ::Tangent) to see if that errors.

[Red-Portal]

Yeah, this issue just poised everything I was working on. Anybody got around this by rolling back? Not sure whether rolling back Zygote or ChainRules works

[DhairyaLGandhi]

Zygote would need to be rolled back.