Tracker does not work with multilevel model

[sjwild]

Howdy all,

I am trying to get Tracker and/or Zygote to work with a multilevel model, but unfortunately I get an error message when I try. Tracker works fine with a varying intercept model, but not with a varying slope model. Zygote does not seem to work with either.

The code below is a reproducible example (modified slightly to use the sleepstudy dataset). It runs faster using forwarddiff than it does using Tracker, but hopefully it will get my point across.

using Plots, StatsPlots
using LinearAlgebra
using Distributions
using Turing
using Zygote
using RDatasets
using BenchmarkTools
using CSV, DataFrames


# Load the oft-used sleepstudy dataset
df = dataset("lme4", "sleepstudy") |> DataFrame

# prepare groups
subject_map = Dict(key => idx for (idx, key) in enumerate(unique(df.Subject)))
df.Subject_id = [subject_map[sj] for sj in df.Subject]


# calculate mean days for each subject (doesn't matter for this example,
# but does for other datasets))
df.MeanDays = Vector(undef, size(df, 1))
df_group = groupby(df, :Subject)
for i in 1:length(df_group)
    df_group[i][:, :MeanDays] = df_group[i].Days .- mean(df_group[i].Days)
end


# create, y, x, z, and index
y = df.Reaction ./ 100
x = df.MeanDays
subject_id = df.Subject_id


# Set backend
Turing.setadbackend(:tracker)
#Turing.setadbackend(:zygote) # doesn't work

# First model - varying intercept only
# Takes about 30 seconds to run using Tracker
# Takes about 15 seconds on forwarddiff
@model function varying_intercept(y, X, gr_id)

    n_gr = length(unique(gr_id))

    # priors
    σ ~ truncated(Normal(0, 10), 0, Inf)
    τ ~ truncated(Normal(0, 10), 0, Inf)
    α ~ Normal(0, 10)
    β ~ Normal(0, 10)
    μ₁ ~ filldist(Normal(0, τ), n_gr)

    # model
    yhat = α .+ X * β .+ μ₁[gr_id]
    y ~ MvNormal(yhat, σ)

end

model = varying_intercept(y, x, subject_id)
@time chn = sample(model, Turing.NUTS(500, 0.65), 1500)


# varying slopes model. Does not run with Tracker for reverse diff
# Takes about 5 minutes to run using forward diff
@model function varying_slopes(y, X, Z, gr_id)

    # dims 
    #n_x, m_x = size(X)
    #n_z, m_z = size(Z)
    m_z = 1
    n_gr = length(unique(gr_id))

    # priors
    Ρ ~ LKJ((m_z + 1), 2.)
    τ ~ filldist(truncated(Cauchy(0, 2), 0, Inf), (m_z + 1))
    β ~ Normal(0, 5)
    α ~ Normal(0, 5)
    σ ~ truncated(Cauchy(0, 2), 0, Inf)
    z_Rho ~ filldist(Normal(0, 1), (m_z + 1), n_gr)

    # build Rho 
    Ρ = (Ρ' + Ρ) / 2
    L_Ρ = LinearAlgebra.cholesky(Ρ).L
    μ = LinearAlgebra.diagm(τ) * L_Ρ * z_Rho

    μ₁ = μ[1, :]
    μ₂ = μ[2, :]

    # model
    ŷ = α .+ μ₁[gr_id] .+ X .* β .+ Z .* μ₂[gr_id]
    y ~ MvNormal(ŷ, σ)

end

model = varying_slopes(y, x, x, subject_id)
@time chn = sample(model, Turing.NUTS(500, 0.65), 1500)

When I run the code using Tracker for the varying slopes model, I get the following error message:

ERROR: MethodError: no method matching Float64(::Tracker.TrackedReal{Float64})
Closest candidates are:
Float64(::Real, ::RoundingMode) where T<:AbstractFloat at rounding.jl:200
Float64(::T) where T<:Number at boot.jl:716
Float64(::Float32) at float.jl:255
...
Stacktrace:
[1] convert(::Type{Float64}, ::Tracker.TrackedReal{Float64}) at ./number.jl:7
[2] setindex! at ./array.jl:849 [inlined]
[3] setindex! at /Users/julia/buildbot/worker/package_macos64/build/usr/share/julia/stdlib/v1.5/LinearAlgebra/src/triangular.jl:234 [inlined]
[4] (::Bijectors.var"#pullback_link_chol_lkj#220"{UpperTriangular{Float64,Array{Float64,2}},Int64,UpperTriangular{Float64,Array{Float64,2}}})(::TrackedArray{…,Array{Float64,2}}) at /Users/stephenjwild/.julia/packages/Bijectors/L39Ij/src/compat/tracker.jl:426
[5] back_(::Tracker.Grads, ::Tracker.Call{Bijectors.var"#pullback_link_chol_lkj#220"{UpperTriangular{Float64,Array{Float64,2}},Int64,UpperTriangular{Float64,Array{Float64,2}}},Tuple{Tracker.Tracked{UpperTriangular{Float64,Array{Float64,2}}}}}, ::TrackedArray{…,Array{Float64,2}}) at /Users/stephenjwild/.julia/packages/Tracker/Z8Onp/src/back.jl:110
[6] back(::Tracker.Grads, ::Tracker.Tracked{UpperTriangular{Float64,Array{Float64,2}}}, ::TrackedArray{…,Array{Float64,2}}) at /Users/stephenjwild/.julia/packages/Tracker/Z8Onp/src/back.jl:125
[7] (::Tracker.var"#16#17"{Tracker.Grads})(::Tracker.Tracked{UpperTriangular{Float64,Array{Float64,2}}}, ::TrackedArray{…,Array{Float64,2}}) at /Users/stephenjwild/.julia/packages/Tracker/Z8Onp/src/back.jl:113
[8] foreach(::Function, ::Tuple{Tracker.Tracked{UpperTriangular{Float64,Array{Float64,2}}},Nothing}, ::Tuple{TrackedArray{…,Array{Float64,2}},TrackedArray{…,Array{Float64,2}}}) at ./abstractarray.jl:2010
[9] back_(::Tracker.Grads, ::Tracker.Call{Tracker.var"#back#564"{2,typeof(+),Tuple{TrackedArray{…,UpperTriangular{Float64,Array{Float64,2}}},Array{Float64,2}}},Tuple{Tracker.Tracked{UpperTriangular{Float64,Array{Float64,2}}},Nothing}}, ::TrackedArray{…,Array{Float64,2}}) at /Users/stephenjwild/.julia/packages/Tracker/Z8Onp/src/back.jl:113
[10] back(::Tracker.Grads, ::Tracker.Tracked{Array{Float64,2}}, ::TrackedArray{…,Array{Float64,2}}) at /Users/stephenjwild/.julia/packages/Tracker/Z8Onp/src/back.jl:125
[11] #16 at /Users/stephenjwild/.julia/packages/Tracker/Z8Onp/src/back.jl:113 [inlined]
[12] foreach at ./abstractarray.jl:2010 [inlined]
... (the last 4 lines are repeated 1 more time)
[17] back_(::Tracker.Grads, ::Tracker.Call{Tracker.var"#173#174"{Tracker.TrackedReal{Float64}},Tuple{Tracker.Tracked{Float64}}}, ::Tracker.TrackedReal{Float64}) at /Users/stephenjwild/.julia/packages/Tracker/Z8Onp/src/back.jl:113
[18] back(::Tracker.Grads, ::Tracker.Tracked{Float64}, ::Tracker.TrackedReal{Float64}) at /Users/stephenjwild/.julia/packages/Tracker/Z8Onp/src/back.jl:123
[19] #16 at /Users/stephenjwild/.julia/packages/Tracker/Z8Onp/src/back.jl:113 [inlined]
[20] foreach at ./abstractarray.jl:2010 [inlined]
[21] back_(::Tracker.Grads, ::Tracker.Call{Tracker.var"#281#284"{Int64},Tuple{Nothing,Tracker.Tracked{Float64}}}, ::Tracker.TrackedReal{Float64}) at /Users/stephenjwild/.julia/packages/Tracker/Z8Onp/src/back.jl:113
[22] back(::Tracker.Grads, ::Tracker.Tracked{Float64}, ::Tracker.TrackedReal{Float64}) at /Users/stephenjwild/.julia/packages/Tracker/Z8Onp/src/back.jl:125
... (the last 4 lines are repeated 15 more times)
[83] #18 at /Users/stephenjwild/.julia/packages/Tracker/Z8Onp/src/back.jl:140 [inlined]
[84] (::Tracker.var"#21#23"{Tracker.var"#18#19"{Tracker.Params,Tracker.TrackedReal{Float64}}})(::Int64) at /Users/stephenjwild/.julia/packages/Tracker/Z8Onp/src/back.jl:149
[85] gradient_logp(::Turing.Core.TrackerAD, ::Array{Float64,1}, ::DynamicPPL.VarInfo{NamedTuple{(:Ρ, :τ, :β, :α, :σ, :z_Rho),Tuple{DynamicPPL.Metadata{Dict{DynamicPPL.VarName{:Ρ,Tuple{}},Int64},Array{LKJ{Float64,Int64},1},Array{DynamicPPL.VarName{:Ρ,Tuple{}},1},Array{Float64,1},Array{Set{DynamicPPL.Selector},1}},DynamicPPL.Metadata{Dict{DynamicPPL.VarName{:τ,Tuple{}},Int64},Array{Product{Continuous,Truncated{Cauchy{Float64},Continuous,Float64},FillArrays.Fill{Truncated{Cauchy{Float64},Continuous,Float64},1,Tuple{Base.OneTo{Int64}}}},1},Array{DynamicPPL.VarName{:τ,Tuple{}},1},Array{Float64,1},Array{Set{DynamicPPL.Selector},1}},DynamicPPL.Metadata{Dict{DynamicPPL.VarName{:β,Tuple{}},Int64},Array{Normal{Float64},1},Array{DynamicPPL.VarName{:β,Tuple{}},1},Array{Float64,1},Array{Set{DynamicPPL.Selector},1}},DynamicPPL.Metadata{Dict{DynamicPPL.VarName{:α,Tuple{}},Int64},Array{Normal{Float64},1},Array{DynamicPPL.VarName{:α,Tuple{}},1},Array{Float64,1},Array{Set{DynamicPPL.Selector},1}},DynamicPPL.Metadata{Dict{DynamicPPL.VarName{:σ,Tuple{}},Int64},Array{Truncated{Cauchy{Float64},Continuous,Float64},1},Array{DynamicPPL.VarName{:σ,Tuple{}},1},Array{Float64,1},Array{Set{DynamicPPL.Selector},1}},DynamicPPL.Metadata{Dict{DynamicPPL.VarName{:z_Rho,Tuple{}},Int64},Array{DistributionsAD.MatrixOfUnivariate{Continuous,Normal{Float64},FillArrays.Fill{Normal{Float64},2,Tuple{Base.OneTo{Int64},Base.OneTo{Int64}}}},1},Array{DynamicPPL.VarName{:z_Rho,Tuple{}},1},Array{Float64,1},Array{Set{DynamicPPL.Selector},1}}}},Float64}, ::DynamicPPL.Model{var"#11#12",(:y, :X, :Z, :gr_id),(),(),Tuple{Array{Float64,1},Array{Any,1},Array{Any,1},Array{Int64,1}},Tuple{}}, ::DynamicPPL.Sampler{NUTS{Turing.Core.TrackerAD,(),AdvancedHMC.DiagEuclideanMetric}}, ::DynamicPPL.DefaultContext) at /Users/stephenjwild/.julia/packages/Turing/XLLTf/src/core/ad.jl:147
[86] gradient_logp at /Users/stephenjwild/.julia/packages/Turing/XLLTf/src/core/ad.jl:83 [inlined] (repeats 2 times)
[87] ∂logπ∂θ at /Users/stephenjwild/.julia/packages/Turing/XLLTf/src/inference/hmc.jl:433 [inlined]
[88] ∂H∂θ at /Users/stephenjwild/.julia/packages/AdvancedHMC/MIxdK/src/hamiltonian.jl:31 [inlined]
[89] phasepoint at /Users/stephenjwild/.julia/packages/AdvancedHMC/MIxdK/src/hamiltonian.jl:69 [inlined]
[90] phasepoint(::Random._GLOBAL_RNG, ::Array{Float64,1}, ::AdvancedHMC.Hamiltonian{AdvancedHMC.DiagEuclideanMetric{Float64,Array{Float64,1}},Turing.Inference.var"#logπ#52"{DynamicPPL.VarInfo{NamedTuple{(:Ρ, :τ, :β, :α, :σ, :z_Rho),Tuple{DynamicPPL.Metadata{Dict{DynamicPPL.VarName{:Ρ,Tuple{}},Int64},Array{LKJ{Float64,Int64},1},Array{DynamicPPL.VarName{:Ρ,Tuple{}},1},Array{Float64,1},Array{Set{DynamicPPL.Selector},1}},DynamicPPL.Metadata{Dict{DynamicPPL.VarName{:τ,Tuple{}},Int64},Array{Product{Continuous,Truncated{Cauchy{Float64},Continuous,Float64},FillArrays.Fill{Truncated{Cauchy{Float64},Continuous,Float64},1,Tuple{Base.OneTo{Int64}}}},1},Array{DynamicPPL.VarName{:τ,Tuple{}},1},Array{Float64,1},Array{Set{DynamicPPL.Selector},1}},DynamicPPL.Metadata{Dict{DynamicPPL.VarName{:β,Tuple{}},Int64},Array{Normal{Float64},1},Array{DynamicPPL.VarName{:β,Tuple{}},1},Array{Float64,1},Array{Set{DynamicPPL.Selector},1}},DynamicPPL.Metadata{Dict{DynamicPPL.VarName{:α,Tuple{}},Int64},Array{Normal{Float64},1},Array{DynamicPPL.VarName{:α,Tuple{}},1},Array{Float64,1},Array{Set{DynamicPPL.Selector},1}},DynamicPPL.Metadata{Dict{DynamicPPL.VarName{:σ,Tuple{}},Int64},Array{Truncated{Cauchy{Float64},Continuous,Float64},1},Array{DynamicPPL.VarName{:σ,Tuple{}},1},Array{Float64,1},Array{Set{DynamicPPL.Selector},1}},DynamicPPL.Metadata{Dict{DynamicPPL.VarName{:z_Rho,Tuple{}},Int64},Array{DistributionsAD.MatrixOfUnivariate{Continuous,Normal{Float64},FillArrays.Fill{Normal{Float64},2,Tuple{Base.OneTo{Int64},Base.OneTo{Int64}}}},1},Array{DynamicPPL.VarName{:z_Rho,Tuple{}},1},Array{Float64,1},Array{Set{DynamicPPL.Selector},1}}}},Float64},DynamicPPL.Sampler{NUTS{Turing.Core.TrackerAD,(),AdvancedHMC.DiagEuclideanMetric}},DynamicPPL.Model{var"#11#12",(:y, :X, :Z, :gr_id),(),(),Tuple{Array{Float64,1},Array{Any,1},Array{Any,1},Array{Int64,1}},Tuple{}}},Turing.Inference.var"#∂logπ∂θ#51"{DynamicPPL.VarInfo{NamedTuple{(:Ρ, :τ, :β, :α, :σ, :z_Rho),Tuple{DynamicPPL.Metadata{Dict{DynamicPPL.VarName{:Ρ,Tuple{}},Int64},Array{LKJ{Float64,Int64},1},Array{DynamicPPL.VarName{:Ρ,Tuple{}},1},Array{Float64,1},Array{Set{DynamicPPL.Selector},1}},DynamicPPL.Metadata{Dict{DynamicPPL.VarName{:τ,Tuple{}},Int64},Array{Product{Continuous,Truncated{Cauchy{Float64},Continuous,Float64},FillArrays.Fill{Truncated{Cauchy{Float64},Continuous,Float64},1,Tuple{Base.OneTo{Int64}}}},1},Array{DynamicPPL.VarName{:τ,Tuple{}},1},Array{Float64,1},Array{Set{DynamicPPL.Selector},1}},DynamicPPL.Metadata{Dict{DynamicPPL.VarName{:β,Tuple{}},Int64},Array{Normal{Float64},1},Array{DynamicPPL.VarName{:β,Tuple{}},1},Array{Float64,1},Array{Set{DynamicPPL.Selector},1}},DynamicPPL.Metadata{Dict{DynamicPPL.VarName{:α,Tuple{}},Int64},Array{Normal{Float64},1},Array{DynamicPPL.VarName{:α,Tuple{}},1},Array{Float64,1},Array{Set{DynamicPPL.Selector},1}},DynamicPPL.Metadata{Dict{DynamicPPL.VarName{:σ,Tuple{}},Int64},Array{Truncated{Cauchy{Float64},Continuous,Float64},1},Array{DynamicPPL.VarName{:σ,Tuple{}},1},Array{Float64,1},Array{Set{DynamicPPL.Selector},1}},DynamicPPL.Metadata{Dict{DynamicPPL.VarName{:z_Rho,Tuple{}},Int64},Array{DistributionsAD.MatrixOfUnivariate{Continuous,Normal{Float64},FillArrays.Fill{Normal{Float64},2,Tuple{Base.OneTo{Int64},Base.OneTo{Int64}}}},1},Array{DynamicPPL.VarName{:z_Rho,Tuple{}},1},Array{Float64,1},Array{Set{DynamicPPL.Selector},1}}}},Float64},DynamicPPL.Sampler{NUTS{Turing.Core.TrackerAD,(),AdvancedHMC.DiagEuclideanMetric}},DynamicPPL.Model{var"#11#12",(:y, :X, :Z, :gr_id),(),(),Tuple{Array{Float64,1},Array{Any,1},Array{Any,1},Array{Int64,1}},Tuple{}}}}) at /Users/stephenjwild/.julia/packages/AdvancedHMC/MIxdK/src/hamiltonian.jl:139
[91] initialstep(::Random._GLOBAL_RNG, ::DynamicPPL.Model{var"#11#12",(:y, :X, :Z, :gr_id),(),(),Tuple{Array{Float64,1},Array{Any,1},Array{Any,1},Array{Int64,1}},Tuple{}}, ::DynamicPPL.Sampler{NUTS{Turing.Core.TrackerAD,(),AdvancedHMC.DiagEuclideanMetric}}, ::DynamicPPL.VarInfo{NamedTuple{(:Ρ, :τ, :β, :α, :σ, :z_Rho),Tuple{DynamicPPL.Metadata{Dict{DynamicPPL.VarName{:Ρ,Tuple{}},Int64},Array{LKJ{Float64,Int64},1},Array{DynamicPPL.VarName{:Ρ,Tuple{}},1},Array{Float64,1},Array{Set{DynamicPPL.Selector},1}},DynamicPPL.Metadata{Dict{DynamicPPL.VarName{:τ,Tuple{}},Int64},Array{Product{Continuous,Truncated{Cauchy{Float64},Continuous,Float64},FillArrays.Fill{Truncated{Cauchy{Float64},Continuous,Float64},1,Tuple{Base.OneTo{Int64}}}},1},Array{DynamicPPL.VarName{:τ,Tuple{}},1},Array{Float64,1},Array{Set{DynamicPPL.Selector},1}},DynamicPPL.Metadata{Dict{DynamicPPL.VarName{:β,Tuple{}},Int64},Array{Normal{Float64},1},Array{DynamicPPL.VarName{:β,Tuple{}},1},Array{Float64,1},Array{Set{DynamicPPL.Selector},1}},DynamicPPL.Metadata{Dict{DynamicPPL.VarName{:α,Tuple{}},Int64},Array{Normal{Float64},1},Array{DynamicPPL.VarName{:α,Tuple{}},1},Array{Float64,1},Array{Set{DynamicPPL.Selector},1}},DynamicPPL.Metadata{Dict{DynamicPPL.VarName{:σ,Tuple{}},Int64},Array{Truncated{Cauchy{Float64},Continuous,Float64},1},Array{DynamicPPL.VarName{:σ,Tuple{}},1},Array{Float64,1},Array{Set{DynamicPPL.Selector},1}},DynamicPPL.Metadata{Dict{DynamicPPL.VarName{:z_Rho,Tuple{}},Int64},Array{DistributionsAD.MatrixOfUnivariate{Continuous,Normal{Float64},FillArrays.Fill{Normal{Float64},2,Tuple{Base.OneTo{Int64},Base.OneTo{Int64}}}},1},Array{DynamicPPL.VarName{:z_Rho,Tuple{}},1},Array{Float64,1},Array{Set{DynamicPPL.Selector},1}}}},Float64}; init_params::Nothing, nadapts::Int64, kwargs::Base.Iterators.Pairs{Union{},Union{},Tuple{},NamedTuple{(),Tuple{}}}) at /Users/stephenjwild/.julia/packages/Turing/XLLTf/src/inference/hmc.jl:167
[92] step(::Random._GLOBAL_RNG, ::DynamicPPL.Model{var"#11#12",(:y, :X, :Z, :gr_id),(),(),Tuple{Array{Float64,1},Array{Any,1},Array{Any,1},Array{Int64,1}},Tuple{}}, ::DynamicPPL.Sampler{NUTS{Turing.Core.TrackerAD,(),AdvancedHMC.DiagEuclideanMetric}}; resume_from::Nothing, kwargs::Base.Iterators.Pairs{Symbol,Int64,Tuple{Symbol},NamedTuple{(:nadapts,),Tuple{Int64}}}) at /Users/stephenjwild/.julia/packages/DynamicPPL/wf0dU/src/sampler.jl:83
[93] macro expansion at /Users/stephenjwild/.julia/packages/AbstractMCMC/Nw3Wn/src/sample.jl:78 [inlined]
[94] macro expansion at /Users/stephenjwild/.julia/packages/ProgressLogging/6KXlp/src/ProgressLogging.jl:328 [inlined]
[95] macro expansion at /Users/stephenjwild/.julia/packages/AbstractMCMC/Nw3Wn/src/logging.jl:8 [inlined]
[96] mcmcsample(::Random._GLOBAL_RNG, ::DynamicPPL.Model{var"#11#12",(:y, :X, :Z, :gr_id),(),(),Tuple{Array{Float64,1},Array{Any,1},Array{Any,1},Array{Int64,1}},Tuple{}}, ::DynamicPPL.Sampler{NUTS{Turing.Core.TrackerAD,(),AdvancedHMC.DiagEuclideanMetric}}, ::Int64; progress::Bool, progressname::String, callback::Nothing, discard_initial::Int64, thinning::Int64, chain_type::Type{T} where T, kwargs::Base.Iterators.Pairs{Symbol,Int64,Tuple{Symbol},NamedTuple{(:nadapts,),Tuple{Int64}}}) at /Users/stephenjwild/.julia/packages/AbstractMCMC/Nw3Wn/src/sample.jl:76
[97] sample(::Random._GLOBAL_RNG, ::DynamicPPL.Model{var"#11#12",(:y, :X, :Z, :gr_id),(),(),Tuple{Array{Float64,1},Array{Any,1},Array{Any,1},Array{Int64,1}},Tuple{}}, ::DynamicPPL.Sampler{NUTS{Turing.Core.TrackerAD,(),AdvancedHMC.DiagEuclideanMetric}}, ::Int64; chain_type::Type{T} where T, resume_from::Nothing, progress::Bool, nadapts::Int64, discard_adapt::Bool, discard_initial::Int64, kwargs::Base.Iterators.Pairs{Union{},Union{},Tuple{},NamedTuple{(),Tuple{}}}) at /Users/stephenjwild/.julia/packages/Turing/XLLTf/src/inference/hmc.jl:133
[98] sample at /Users/stephenjwild/.julia/packages/Turing/XLLTf/src/inference/hmc.jl:116 [inlined]
[99] #sample#2 at /Users/stephenjwild/.julia/packages/Turing/XLLTf/src/inference/Inference.jl:141 [inlined]
[100] sample at /Users/stephenjwild/.julia/packages/Turing/XLLTf/src/inference/Inference.jl:141 [inlined]
[101] #sample#1 at /Users/stephenjwild/.julia/packages/Turing/XLLTf/src/inference/Inference.jl:131 [inlined]
[102] sample(::DynamicPPL.Model{var"#11#12",(:y, :X, :Z, :gr_id),(),(),Tuple{Array{Float64,1},Array{Any,1},Array{Any,1},Array{Int64,1}},Tuple{}}, ::NUTS{Turing.Core.TrackerAD,(),AdvancedHMC.DiagEuclideanMetric}, ::Int64) at /Users/stephenjwild/.julia/packages/Turing/XLLTf/src/inference/Inference.jl:131
[103] top-level scope at ./timing.jl:174 [inlined]

If I try Zygote with either model, I get the following error:

ERROR: Mutating arrays is not supported
Stacktrace:
[1] error(::String) at ./error.jl:33
[2] (::Zygote.var"#372#373")(::Nothing) at /Users/stephenjwild/.julia/packages/Zygote/ggM8Z/src/lib/array.jl:65
[3] (::Zygote.var"#2265#back#374"{Zygote.var"#372#373"})(::Nothing) at /Users/stephenjwild/.julia/packages/ZygoteRules/OjfTt/src/adjoint.jl:59
[4] unique_from at ./set.jl:157 [inlined]
[5] (::typeof(∂(unique_from)))(::Nothing) at /Users/stephenjwild/.julia/packages/Zygote/ggM8Z/src/compiler/interface2.jl:0
[6] unique at ./set.jl:135 [inlined]
[7] (::typeof(∂(invoke)))(::Nothing) at /Users/stephenjwild/.julia/packages/Zygote/ggM8Z/src/compiler/interface2.jl:0
[8] _unique_dims at ./multidimensional.jl:1474 [inlined]
[9] #unique#450 at ./multidimensional.jl:1472 [inlined]
[10] unique at ./multidimensional.jl:1472 [inlined]
[11] (::typeof(∂(unique)))(::Nothing) at /Users/stephenjwild/.julia/packages/Zygote/ggM8Z/src/compiler/interface2.jl:0
[12] #17 at ./REPL[34]:2 [inlined]
[13] (::typeof(∂(#17)))(::Nothing) at /Users/stephenjwild/.julia/packages/Zygote/ggM8Z/src/compiler/interface2.jl:0
[14] macro expansion at /Users/stephenjwild/.julia/packages/DynamicPPL/wf0dU/src/model.jl:0 [inlined]
[15] _evaluate at /Users/stephenjwild/.julia/packages/DynamicPPL/wf0dU/src/model.jl:154 [inlined]
[16] (::typeof(∂(_evaluate)))(::Nothing) at /Users/stephenjwild/.julia/packages/Zygote/ggM8Z/src/compiler/interface2.jl:0
[17] evaluate_threadunsafe at /Users/stephenjwild/.julia/packages/DynamicPPL/wf0dU/src/model.jl:127 [inlined]
[18] (::typeof(∂(evaluate_threadunsafe)))(::Nothing) at /Users/stephenjwild/.julia/packages/Zygote/ggM8Z/src/compiler/interface2.jl:0
[19] Model at /Users/stephenjwild/.julia/packages/DynamicPPL/wf0dU/src/model.jl:92 [inlined]
[20] (::typeof(∂(λ)))(::Nothing) at /Users/stephenjwild/.julia/packages/Zygote/ggM8Z/src/compiler/interface2.jl:0
[21] #150 at /Users/stephenjwild/.julia/packages/Zygote/ggM8Z/src/lib/lib.jl:191 [inlined]
[22] #1693#back at /Users/stephenjwild/.julia/packages/ZygoteRules/OjfTt/src/adjoint.jl:59 [inlined]
[23] Model at /Users/stephenjwild/.julia/packages/DynamicPPL/wf0dU/src/model.jl:98 [inlined]
[24] f at /Users/stephenjwild/.julia/packages/Turing/XLLTf/src/core/ad.jl:165 [inlined]
[25] (::typeof(∂(λ)))(::Int64) at /Users/stephenjwild/.julia/packages/Zygote/ggM8Z/src/compiler/interface2.jl:0
[26] (::Zygote.var"#41#42"{typeof(∂(λ))})(::Int64) at /Users/stephenjwild/.julia/packages/Zygote/ggM8Z/src/compiler/interface.jl:40
[27] gradient_logp(::Turing.Core.ZygoteAD, ::Array{Float64,1}, ::DynamicPPL.VarInfo{NamedTuple{(:σ, :τ, :α, :β, :μ₁),Tuple{DynamicPPL.Metadata{Dict{DynamicPPL.VarName{:σ,Tuple{}},Int64},Array{Truncated{Normal{Float64},Continuous,Float64},1},Array{DynamicPPL.VarName{:σ,Tuple{}},1},Array{Float64,1},Array{Set{DynamicPPL.Selector},1}},DynamicPPL.Metadata{Dict{DynamicPPL.VarName{:τ,Tuple{}},Int64},Array{Truncated{Normal{Float64},Continuous,Float64},1},Array{DynamicPPL.VarName{:τ,Tuple{}},1},Array{Float64,1},Array{Set{DynamicPPL.Selector},1}},DynamicPPL.Metadata{Dict{DynamicPPL.VarName{:α,Tuple{}},Int64},Array{Normal{Float64},1},Array{DynamicPPL.VarName{:α,Tuple{}},1},Array{Float64,1},Array{Set{DynamicPPL.Selector},1}},DynamicPPL.Metadata{Dict{DynamicPPL.VarName{:β,Tuple{}},Int64},Array{Normal{Float64},1},Array{DynamicPPL.VarName{:β,Tuple{}},1},Array{Float64,1},Array{Set{DynamicPPL.Selector},1}},DynamicPPL.Metadata{Dict{DynamicPPL.VarName{:μ₁,Tuple{}},Int64},Array{DistributionsAD.TuringScalMvNormal{Array{Float64,1},Float64},1},Array{DynamicPPL.VarName{:μ₁,Tuple{}},1},Array{Float64,1},Array{Set{DynamicPPL.Selector},1}}}},Float64}, ::DynamicPPL.Model{var"#17#18",(:y, :X, :gr_id),(),(),Tuple{Array{Float64,1},Array{Any,1},Array{Int64,1}},Tuple{}}, ::DynamicPPL.Sampler{NUTS{Turing.Core.ZygoteAD,(),AdvancedHMC.DiagEuclideanMetric}}, ::DynamicPPL.DefaultContext) at /Users/stephenjwild/.julia/packages/Turing/XLLTf/src/core/ad.jl:171
[28] gradient_logp at /Users/stephenjwild/.julia/packages/Turing/XLLTf/src/core/ad.jl:83 [inlined] (repeats 2 times)
[29] ∂logπ∂θ at /Users/stephenjwild/.julia/packages/Turing/XLLTf/src/inference/hmc.jl:433 [inlined]
[30] ∂H∂θ at /Users/stephenjwild/.julia/packages/AdvancedHMC/MIxdK/src/hamiltonian.jl:31 [inlined]
[31] phasepoint at /Users/stephenjwild/.julia/packages/AdvancedHMC/MIxdK/src/hamiltonian.jl:69 [inlined]
[32] phasepoint(::Random._GLOBAL_RNG, ::Array{Float64,1}, ::AdvancedHMC.Hamiltonian{AdvancedHMC.DiagEuclideanMetric{Float64,Array{Float64,1}},Turing.Inference.var"#logπ#52"{DynamicPPL.VarInfo{NamedTuple{(:σ, :τ, :α, :β, :μ₁),Tuple{DynamicPPL.Metadata{Dict{DynamicPPL.VarName{:σ,Tuple{}},Int64},Array{Truncated{Normal{Float64},Continuous,Float64},1},Array{DynamicPPL.VarName{:σ,Tuple{}},1},Array{Float64,1},Array{Set{DynamicPPL.Selector},1}},DynamicPPL.Metadata{Dict{DynamicPPL.VarName{:τ,Tuple{}},Int64},Array{Truncated{Normal{Float64},Continuous,Float64},1},Array{DynamicPPL.VarName{:τ,Tuple{}},1},Array{Float64,1},Array{Set{DynamicPPL.Selector},1}},DynamicPPL.Metadata{Dict{DynamicPPL.VarName{:α,Tuple{}},Int64},Array{Normal{Float64},1},Array{DynamicPPL.VarName{:α,Tuple{}},1},Array{Float64,1},Array{Set{DynamicPPL.Selector},1}},DynamicPPL.Metadata{Dict{DynamicPPL.VarName{:β,Tuple{}},Int64},Array{Normal{Float64},1},Array{DynamicPPL.VarName{:β,Tuple{}},1},Array{Float64,1},Array{Set{DynamicPPL.Selector},1}},DynamicPPL.Metadata{Dict{DynamicPPL.VarName{:μ₁,Tuple{}},Int64},Array{DistributionsAD.TuringScalMvNormal{Array{Float64,1},Float64},1},Array{DynamicPPL.VarName{:μ₁,Tuple{}},1},Array{Float64,1},Array{Set{DynamicPPL.Selector},1}}}},Float64},DynamicPPL.Sampler{NUTS{Turing.Core.ZygoteAD,(),AdvancedHMC.DiagEuclideanMetric}},DynamicPPL.Model{var"#17#18",(:y, :X, :gr_id),(),(),Tuple{Array{Float64,1},Array{Any,1},Array{Int64,1}},Tuple{}}},Turing.Inference.var"#∂logπ∂θ#51"{DynamicPPL.VarInfo{NamedTuple{(:σ, :τ, :α, :β, :μ₁),Tuple{DynamicPPL.Metadata{Dict{DynamicPPL.VarName{:σ,Tuple{}},Int64},Array{Truncated{Normal{Float64},Continuous,Float64},1},Array{DynamicPPL.VarName{:σ,Tuple{}},1},Array{Float64,1},Array{Set{DynamicPPL.Selector},1}},DynamicPPL.Metadata{Dict{DynamicPPL.VarName{:τ,Tuple{}},Int64},Array{Truncated{Normal{Float64},Continuous,Float64},1},Array{DynamicPPL.VarName{:τ,Tuple{}},1},Array{Float64,1},Array{Set{DynamicPPL.Selector},1}},DynamicPPL.Metadata{Dict{DynamicPPL.VarName{:α,Tuple{}},Int64},Array{Normal{Float64},1},Array{DynamicPPL.VarName{:α,Tuple{}},1},Array{Float64,1},Array{Set{DynamicPPL.Selector},1}},DynamicPPL.Metadata{Dict{DynamicPPL.VarName{:β,Tuple{}},Int64},Array{Normal{Float64},1},Array{DynamicPPL.VarName{:β,Tuple{}},1},Array{Float64,1},Array{Set{DynamicPPL.Selector},1}},DynamicPPL.Metadata{Dict{DynamicPPL.VarName{:μ₁,Tuple{}},Int64},Array{DistributionsAD.TuringScalMvNormal{Array{Float64,1},Float64},1},Array{DynamicPPL.VarName{:μ₁,Tuple{}},1},Array{Float64,1},Array{Set{DynamicPPL.Selector},1}}}},Float64},DynamicPPL.Sampler{NUTS{Turing.Core.ZygoteAD,(),AdvancedHMC.DiagEuclideanMetric}},DynamicPPL.Model{var"#17#18",(:y, :X, :gr_id),(),(),Tuple{Array{Float64,1},Array{Any,1},Array{Int64,1}},Tuple{}}}}) at /Users/stephenjwild/.julia/packages/AdvancedHMC/MIxdK/src/hamiltonian.jl:139
[33] initialstep(::Random._GLOBAL_RNG, ::DynamicPPL.Model{var"#17#18",(:y, :X, :gr_id),(),(),Tuple{Array{Float64,1},Array{Any,1},Array{Int64,1}},Tuple{}}, ::DynamicPPL.Sampler{NUTS{Turing.Core.ZygoteAD,(),AdvancedHMC.DiagEuclideanMetric}}, ::DynamicPPL.VarInfo{NamedTuple{(:σ, :τ, :α, :β, :μ₁),Tuple{DynamicPPL.Metadata{Dict{DynamicPPL.VarName{:σ,Tuple{}},Int64},Array{Truncated{Normal{Float64},Continuous,Float64},1},Array{DynamicPPL.VarName{:σ,Tuple{}},1},Array{Float64,1},Array{Set{DynamicPPL.Selector},1}},DynamicPPL.Metadata{Dict{DynamicPPL.VarName{:τ,Tuple{}},Int64},Array{Truncated{Normal{Float64},Continuous,Float64},1},Array{DynamicPPL.VarName{:τ,Tuple{}},1},Array{Float64,1},Array{Set{DynamicPPL.Selector},1}},DynamicPPL.Metadata{Dict{DynamicPPL.VarName{:α,Tuple{}},Int64},Array{Normal{Float64},1},Array{DynamicPPL.VarName{:α,Tuple{}},1},Array{Float64,1},Array{Set{DynamicPPL.Selector},1}},DynamicPPL.Metadata{Dict{DynamicPPL.VarName{:β,Tuple{}},Int64},Array{Normal{Float64},1},Array{DynamicPPL.VarName{:β,Tuple{}},1},Array{Float64,1},Array{Set{DynamicPPL.Selector},1}},DynamicPPL.Metadata{Dict{DynamicPPL.VarName{:μ₁,Tuple{}},Int64},Array{DistributionsAD.TuringScalMvNormal{Array{Float64,1},Float64},1},Array{DynamicPPL.VarName{:μ₁,Tuple{}},1},Array{Float64,1},Array{Set{DynamicPPL.Selector},1}}}},Float64}; init_params::Nothing, nadapts::Int64, kwargs::Base.Iterators.Pairs{Union{},Union{},Tuple{},NamedTuple{(),Tuple{}}}) at /Users/stephenjwild/.julia/packages/Turing/XLLTf/src/inference/hmc.jl:167
[34] step(::Random._GLOBAL_RNG, ::DynamicPPL.Model{var"#17#18",(:y, :X, :gr_id),(),(),Tuple{Array{Float64,1},Array{Any,1},Array{Int64,1}},Tuple{}}, ::DynamicPPL.Sampler{NUTS{Turing.Core.ZygoteAD,(),AdvancedHMC.DiagEuclideanMetric}}; resume_from::Nothing, kwargs::Base.Iterators.Pairs{Symbol,Int64,Tuple{Symbol},NamedTuple{(:nadapts,),Tuple{Int64}}}) at /Users/stephenjwild/.julia/packages/DynamicPPL/wf0dU/src/sampler.jl:83
[35] macro expansion at /Users/stephenjwild/.julia/packages/AbstractMCMC/Nw3Wn/src/sample.jl:78 [inlined]
[36] macro expansion at /Users/stephenjwild/.julia/packages/ProgressLogging/6KXlp/src/ProgressLogging.jl:328 [inlined]
[37] macro expansion at /Users/stephenjwild/.julia/packages/AbstractMCMC/Nw3Wn/src/logging.jl:8 [inlined]
[38] mcmcsample(::Random._GLOBAL_RNG, ::DynamicPPL.Model{var"#17#18",(:y, :X, :gr_id),(),(),Tuple{Array{Float64,1},Array{Any,1},Array{Int64,1}},Tuple{}}, ::DynamicPPL.Sampler{NUTS{Turing.Core.ZygoteAD,(),AdvancedHMC.DiagEuclideanMetric}}, ::Int64; progress::Bool, progressname::String, callback::Nothing, discard_initial::Int64, thinning::Int64, chain_type::Type{T} where T, kwargs::Base.Iterators.Pairs{Symbol,Int64,Tuple{Symbol},NamedTuple{(:nadapts,),Tuple{Int64}}}) at /Users/stephenjwild/.julia/packages/AbstractMCMC/Nw3Wn/src/sample.jl:76
[39] sample(::Random._GLOBAL_RNG, ::DynamicPPL.Model{var"#17#18",(:y, :X, :gr_id),(),(),Tuple{Array{Float64,1},Array{Any,1},Array{Int64,1}},Tuple{}}, ::DynamicPPL.Sampler{NUTS{Turing.Core.ZygoteAD,(),AdvancedHMC.DiagEuclideanMetric}}, ::Int64; chain_type::Type{T} where T, resume_from::Nothing, progress::Bool, nadapts::Int64, discard_adapt::Bool, discard_initial::Int64, kwargs::Base.Iterators.Pairs{Union{},Union{},Tuple{},NamedTuple{(),Tuple{}}}) at /Users/stephenjwild/.julia/packages/Turing/XLLTf/src/inference/hmc.jl:133
[40] sample at /Users/stephenjwild/.julia/packages/Turing/XLLTf/src/inference/hmc.jl:116 [inlined]
[41] #sample#2 at /Users/stephenjwild/.julia/packages/Turing/XLLTf/src/inference/Inference.jl:141 [inlined]
[42] sample at /Users/stephenjwild/.julia/packages/Turing/XLLTf/src/inference/Inference.jl:141 [inlined]
[43] #sample#1 at /Users/stephenjwild/.julia/packages/Turing/XLLTf/src/inference/Inference.jl:131 [inlined]
[44] sample(::DynamicPPL.Model{var"#17#18",(:y, :X, :gr_id),(),(),Tuple{Array{Float64,1},Array{Any,1},Array{Int64,1}},Tuple{}}, ::NUTS{Turing.Core.ZygoteAD,(),AdvancedHMC.DiagEuclideanMetric}, ::Int64) at /Users/stephenjwild/.julia/packages/Turing/XLLTf/src/inference/Inference.jl:131
[45] top-level scope at ./timing.jl:174 [inlined]
[46] top-level scope at ./REPL[42]:0

I'm not sure if it is a bug, my poor coding, or both. Hopefully it's not a bug.

Any advice is welcome.

[devmotion]

The problem is the definition of the LKJ bijector. More specifically, in your example https://github.com/TuringLang/Bijectors.jl/blob/50bb27661d724976c0d6c008cf164d4d35515f9e/src/compat/tracker.jl#L424 creates a mattix of Float64 but the algorithm tries to assign TrackedReal to its elements. This causes the error.

[devmotion]

The use of unique seems to be problematic with Zygote (I assume, more generally and not only in your example), and therefore both models don't work with it.

[devmotion]

Generally, it seems inefficient to compute n_gr inside the model - it will be recomputed every time you sample from the model and every time you compute gradients of the log pdf.

You could make it an additional argument of the models (possibly with a default value of length(unique(groups))). Then the value of n_gr will be fixed when you instantiate your models.

[sjwild]

Thanks for the quick response and advice about n_gr. If remove n_gr from the varying intercept model, it cuts the run time in about half for both Tracker and forwarddiff. Once I figure this out a little more, I'll have to to see the time difference on a larger dataset.

That explains the issue with Tracker. I assume there's not much I can do at my end to keep the LKJ prior and still use Tracker as the backend, correct (at least with the way the model is currently coded)?

As for Zygote, it works for the varying intercept model when I move out unique(length(gr_id)). It does, however, take a lot longer than Tracker.

Unfortunately, Zygote still does not work for the varying slope models. This time I get a message about NamedTuple.

julia> model = varying_slopes(y, x, x, subject_id, n_gr, m_z)
DynamicPPL.Model{var"#25#26",(:y, :X, :Z, :gr_id, :n_gr, :m_z),(),(),Tuple{Array{Float64,1},Array{Any,1},Array{Any,1},Array{Int64,1},Int64,Int64},Tuple{}}(:varying_slopes, var"#25#26"(), (y = [2.4956, 2.587047, 2.508006, 3.214398, 3.568519, 4.146901, 3.822038, 2.901486, 4.305853, 4.6635349999999995 … 2.694117, 2.73474, 2.975968, 3.106316, 2.871726, 3.296076, 3.344818, 3.4321989999999998, 3.691417, 3.641236], X = Any[-4.5, -3.5, -2.5, -1.5, -0.5, 0.5, 1.5, 2.5, 3.5, 4.5 … -4.5, -3.5, -2.5, -1.5, -0.5, 0.5, 1.5, 2.5, 3.5, 4.5], Z = Any[-4.5, -3.5, -2.5, -1.5, -0.5, 0.5, 1.5, 2.5, 3.5, 4.5 … -4.5, -3.5, -2.5, -1.5, -0.5, 0.5, 1.5, 2.5, 3.5, 4.5], gr_id = [1, 1, 1, 1, 1, 1, 1, 1, 1, 1 … 18, 18, 18, 18, 18, 18, 18, 18, 18, 18], n_gr = 18, m_z = 2), NamedTuple())

julia> @time chn = sample(model, Turing.NUTS(500, 0.65), 1500)
ERROR: type NamedTuple has no field first
Stacktrace:
[1] getproperty(::NamedTuple{(:second,),Tuple{Array{Float64,1}}}, ::Symbol) at ./Base.jl:33
[2] macro expansion at /Users/stephenjwild/.julia/packages/Zygote/ggM8Z/src/lib/lib.jl:287 [inlined]
[3] (::Zygote.Jnew{Pair{Int64,Array{Float64,1}},Nothing,false})(::NamedTuple{(:second,),Tuple{Array{Float64,1}}}) at /Users/stephenjwild/.julia/packages/Zygote/ggM8Z/src/lib/lib.jl:279
[4] (::Zygote.var"#1727#back#165"{Zygote.Jnew{Pair{Int64,Array{Float64,1}},Nothing,false}})(::NamedTuple{(:second,),Tuple{Array{Float64,1}}}) at /Users/stephenjwild/.julia/packages/ZygoteRules/OjfTt/src/adjoint.jl:59
[5] Pair at ./pair.jl:12 [inlined]
[6] Pair at ./pair.jl:15 [inlined]
[7] (::typeof(∂(Pair)))(::NamedTuple{(:second,),Tuple{Array{Float64,1}}}) at /Users/stephenjwild/.julia/packages/Zygote/ggM8Z/src/compiler/interface2.jl:0
[8] diagm at /Users/julia/buildbot/worker/package_macos64/build/usr/share/julia/stdlib/v1.5/LinearAlgebra/src/dense.jl:327 [inlined]
[9] (::typeof(∂(diagm)))(::Array{Float64,2}) at /Users/stephenjwild/.julia/packages/Zygote/ggM8Z/src/compiler/interface2.jl:0
[10] #25 at ./REPL[59]:12 [inlined]
[11] (::typeof(∂(#25)))(::Nothing) at /Users/stephenjwild/.julia/packages/Zygote/ggM8Z/src/compiler/interface2.jl:0
[12] macro expansion at /Users/stephenjwild/.julia/packages/DynamicPPL/wf0dU/src/model.jl:0 [inlined]
[13] _evaluate at /Users/stephenjwild/.julia/packages/DynamicPPL/wf0dU/src/model.jl:154 [inlined]
[14] (::typeof(∂(_evaluate)))(::Nothing) at /Users/stephenjwild/.julia/packages/Zygote/ggM8Z/src/compiler/interface2.jl:0
[15] evaluate_threadunsafe at /Users/stephenjwild/.julia/packages/DynamicPPL/wf0dU/src/model.jl:127 [inlined]
[16] (::typeof(∂(evaluate_threadunsafe)))(::Nothing) at /Users/stephenjwild/.julia/packages/Zygote/ggM8Z/src/compiler/interface2.jl:0
[17] Model at /Users/stephenjwild/.julia/packages/DynamicPPL/wf0dU/src/model.jl:92 [inlined]
[18] (::typeof(∂(λ)))(::Nothing) at /Users/stephenjwild/.julia/packages/Zygote/ggM8Z/src/compiler/interface2.jl:0
[19] #150 at /Users/stephenjwild/.julia/packages/Zygote/ggM8Z/src/lib/lib.jl:191 [inlined]
[20] #1693#back at /Users/stephenjwild/.julia/packages/ZygoteRules/OjfTt/src/adjoint.jl:59 [inlined]
[21] Model at /Users/stephenjwild/.julia/packages/DynamicPPL/wf0dU/src/model.jl:98 [inlined]
[22] f at /Users/stephenjwild/.julia/packages/Turing/XLLTf/src/core/ad.jl:165 [inlined]
[23] (::typeof(∂(λ)))(::Int64) at /Users/stephenjwild/.julia/packages/Zygote/ggM8Z/src/compiler/interface2.jl:0
[24] (::Zygote.var"#41#42"{typeof(∂(λ))})(::Int64) at /Users/stephenjwild/.julia/packages/Zygote/ggM8Z/src/compiler/interface.jl:40
[25] gradient_logp(::Turing.Core.ZygoteAD, ::Array{Float64,1}, ::DynamicPPL.VarInfo{NamedTuple{(:Ρ, :τ, :β, :α, :σ, :z_Rho),Tuple{DynamicPPL.Metadata{Dict{DynamicPPL.VarName{:Ρ,Tuple{}},Int64},Array{LKJ{Float64,Int64},1},Array{DynamicPPL.VarName{:Ρ,Tuple{}},1},Array{Float64,1},Array{Set{DynamicPPL.Selector},1}},DynamicPPL.Metadata{Dict{DynamicPPL.VarName{:τ,Tuple{}},Int64},Array{Product{Continuous,Truncated{Cauchy{Float64},Continuous,Float64},FillArrays.Fill{Truncated{Cauchy{Float64},Continuous,Float64},1,Tuple{Base.OneTo{Int64}}}},1},Array{DynamicPPL.VarName{:τ,Tuple{}},1},Array{Float64,1},Array{Set{DynamicPPL.Selector},1}},DynamicPPL.Metadata{Dict{DynamicPPL.VarName{:β,Tuple{}},Int64},Array{Normal{Float64},1},Array{DynamicPPL.VarName{:β,Tuple{}},1},Array{Float64,1},Array{Set{DynamicPPL.Selector},1}},DynamicPPL.Metadata{Dict{DynamicPPL.VarName{:α,Tuple{}},Int64},Array{Normal{Float64},1},Array{DynamicPPL.VarName{:α,Tuple{}},1},Array{Float64,1},Array{Set{DynamicPPL.Selector},1}},DynamicPPL.Metadata{Dict{DynamicPPL.VarName{:σ,Tuple{}},Int64},Array{Truncated{Cauchy{Float64},Continuous,Float64},1},Array{DynamicPPL.VarName{:σ,Tuple{}},1},Array{Float64,1},Array{Set{DynamicPPL.Selector},1}},DynamicPPL.Metadata{Dict{DynamicPPL.VarName{:z_Rho,Tuple{}},Int64},Array{DistributionsAD.MatrixOfUnivariate{Continuous,Normal{Float64},FillArrays.Fill{Normal{Float64},2,Tuple{Base.OneTo{Int64},Base.OneTo{Int64}}}},1},Array{DynamicPPL.VarName{:z_Rho,Tuple{}},1},Array{Float64,1},Array{Set{DynamicPPL.Selector},1}}}},Float64}, ::DynamicPPL.Model{var"#25#26",(:y, :X, :Z, :gr_id, :n_gr, :m_z),(),(),Tuple{Array{Float64,1},Array{Any,1},Array{Any,1},Array{Int64,1},Int64,Int64},Tuple{}}, ::DynamicPPL.Sampler{NUTS{Turing.Core.ZygoteAD,(),AdvancedHMC.DiagEuclideanMetric}}, ::DynamicPPL.DefaultContext) at /Users/stephenjwild/.julia/packages/Turing/XLLTf/src/core/ad.jl:171
[26] gradient_logp at /Users/stephenjwild/.julia/packages/Turing/XLLTf/src/core/ad.jl:83 [inlined] (repeats 2 times)
[27] ∂logπ∂θ at /Users/stephenjwild/.julia/packages/Turing/XLLTf/src/inference/hmc.jl:433 [inlined]
[28] ∂H∂θ at /Users/stephenjwild/.julia/packages/AdvancedHMC/MIxdK/src/hamiltonian.jl:31 [inlined]
[29] phasepoint at /Users/stephenjwild/.julia/packages/AdvancedHMC/MIxdK/src/hamiltonian.jl:69 [inlined]
[30] phasepoint(::Random._GLOBAL_RNG, ::Array{Float64,1}, ::AdvancedHMC.Hamiltonian{AdvancedHMC.DiagEuclideanMetric{Float64,Array{Float64,1}},Turing.Inference.var"#logπ#52"{DynamicPPL.VarInfo{NamedTuple{(:Ρ, :τ, :β, :α, :σ, :z_Rho),Tuple{DynamicPPL.Metadata{Dict{DynamicPPL.VarName{:Ρ,Tuple{}},Int64},Array{LKJ{Float64,Int64},1},Array{DynamicPPL.VarName{:Ρ,Tuple{}},1},Array{Float64,1},Array{Set{DynamicPPL.Selector},1}},DynamicPPL.Metadata{Dict{DynamicPPL.VarName{:τ,Tuple{}},Int64},Array{Product{Continuous,Truncated{Cauchy{Float64},Continuous,Float64},FillArrays.Fill{Truncated{Cauchy{Float64},Continuous,Float64},1,Tuple{Base.OneTo{Int64}}}},1},Array{DynamicPPL.VarName{:τ,Tuple{}},1},Array{Float64,1},Array{Set{DynamicPPL.Selector},1}},DynamicPPL.Metadata{Dict{DynamicPPL.VarName{:β,Tuple{}},Int64},Array{Normal{Float64},1},Array{DynamicPPL.VarName{:β,Tuple{}},1},Array{Float64,1},Array{Set{DynamicPPL.Selector},1}},DynamicPPL.Metadata{Dict{DynamicPPL.VarName{:α,Tuple{}},Int64},Array{Normal{Float64},1},Array{DynamicPPL.VarName{:α,Tuple{}},1},Array{Float64,1},Array{Set{DynamicPPL.Selector},1}},DynamicPPL.Metadata{Dict{DynamicPPL.VarName{:σ,Tuple{}},Int64},Array{Truncated{Cauchy{Float64},Continuous,Float64},1},Array{DynamicPPL.VarName{:σ,Tuple{}},1},Array{Float64,1},Array{Set{DynamicPPL.Selector},1}},DynamicPPL.Metadata{Dict{DynamicPPL.VarName{:z_Rho,Tuple{}},Int64},Array{DistributionsAD.MatrixOfUnivariate{Continuous,Normal{Float64},FillArrays.Fill{Normal{Float64},2,Tuple{Base.OneTo{Int64},Base.OneTo{Int64}}}},1},Array{DynamicPPL.VarName{:z_Rho,Tuple{}},1},Array{Float64,1},Array{Set{DynamicPPL.Selector},1}}}},Float64},DynamicPPL.Sampler{NUTS{Turing.Core.ZygoteAD,(),AdvancedHMC.DiagEuclideanMetric}},DynamicPPL.Model{var"#25#26",(:y, :X, :Z, :gr_id, :n_gr, :m_z),(),(),Tuple{Array{Float64,1},Array{Any,1},Array{Any,1},Array{Int64,1},Int64,Int64},Tuple{}}},Turing.Inference.var"#∂logπ∂θ#51"{DynamicPPL.VarInfo{NamedTuple{(:Ρ, :τ, :β, :α, :σ, :z_Rho),Tuple{DynamicPPL.Metadata{Dict{DynamicPPL.VarName{:Ρ,Tuple{}},Int64},Array{LKJ{Float64,Int64},1},Array{DynamicPPL.VarName{:Ρ,Tuple{}},1},Array{Float64,1},Array{Set{DynamicPPL.Selector},1}},DynamicPPL.Metadata{Dict{DynamicPPL.VarName{:τ,Tuple{}},Int64},Array{Product{Continuous,Truncated{Cauchy{Float64},Continuous,Float64},FillArrays.Fill{Truncated{Cauchy{Float64},Continuous,Float64},1,Tuple{Base.OneTo{Int64}}}},1},Array{DynamicPPL.VarName{:τ,Tuple{}},1},Array{Float64,1},Array{Set{DynamicPPL.Selector},1}},DynamicPPL.Metadata{Dict{DynamicPPL.VarName{:β,Tuple{}},Int64},Array{Normal{Float64},1},Array{DynamicPPL.VarName{:β,Tuple{}},1},Array{Float64,1},Array{Set{DynamicPPL.Selector},1}},DynamicPPL.Metadata{Dict{DynamicPPL.VarName{:α,Tuple{}},Int64},Array{Normal{Float64},1},Array{DynamicPPL.VarName{:α,Tuple{}},1},Array{Float64,1},Array{Set{DynamicPPL.Selector},1}},DynamicPPL.Metadata{Dict{DynamicPPL.VarName{:σ,Tuple{}},Int64},Array{Truncated{Cauchy{Float64},Continuous,Float64},1},Array{DynamicPPL.VarName{:σ,Tuple{}},1},Array{Float64,1},Array{Set{DynamicPPL.Selector},1}},DynamicPPL.Metadata{Dict{DynamicPPL.VarName{:z_Rho,Tuple{}},Int64},Array{DistributionsAD.MatrixOfUnivariate{Continuous,Normal{Float64},FillArrays.Fill{Normal{Float64},2,Tuple{Base.OneTo{Int64},Base.OneTo{Int64}}}},1},Array{DynamicPPL.VarName{:z_Rho,Tuple{}},1},Array{Float64,1},Array{Set{DynamicPPL.Selector},1}}}},Float64},DynamicPPL.Sampler{NUTS{Turing.Core.ZygoteAD,(),AdvancedHMC.DiagEuclideanMetric}},DynamicPPL.Model{var"#25#26",(:y, :X, :Z, :gr_id, :n_gr, :m_z),(),(),Tuple{Array{Float64,1},Array{Any,1},Array{Any,1},Array{Int64,1},Int64,Int64},Tuple{}}}}) at /Users/stephenjwild/.julia/packages/AdvancedHMC/MIxdK/src/hamiltonian.jl:139
[31] initialstep(::Random._GLOBAL_RNG, ::DynamicPPL.Model{var"#25#26",(:y, :X, :Z, :gr_id, :n_gr, :m_z),(),(),Tuple{Array{Float64,1},Array{Any,1},Array{Any,1},Array{Int64,1},Int64,Int64},Tuple{}}, ::DynamicPPL.Sampler{NUTS{Turing.Core.ZygoteAD,(),AdvancedHMC.DiagEuclideanMetric}}, ::DynamicPPL.VarInfo{NamedTuple{(:Ρ, :τ, :β, :α, :σ, :z_Rho),Tuple{DynamicPPL.Metadata{Dict{DynamicPPL.VarName{:Ρ,Tuple{}},Int64},Array{LKJ{Float64,Int64},1},Array{DynamicPPL.VarName{:Ρ,Tuple{}},1},Array{Float64,1},Array{Set{DynamicPPL.Selector},1}},DynamicPPL.Metadata{Dict{DynamicPPL.VarName{:τ,Tuple{}},Int64},Array{Product{Continuous,Truncated{Cauchy{Float64},Continuous,Float64},FillArrays.Fill{Truncated{Cauchy{Float64},Continuous,Float64},1,Tuple{Base.OneTo{Int64}}}},1},Array{DynamicPPL.VarName{:τ,Tuple{}},1},Array{Float64,1},Array{Set{DynamicPPL.Selector},1}},DynamicPPL.Metadata{Dict{DynamicPPL.VarName{:β,Tuple{}},Int64},Array{Normal{Float64},1},Array{DynamicPPL.VarName{:β,Tuple{}},1},Array{Float64,1},Array{Set{DynamicPPL.Selector},1}},DynamicPPL.Metadata{Dict{DynamicPPL.VarName{:α,Tuple{}},Int64},Array{Normal{Float64},1},Array{DynamicPPL.VarName{:α,Tuple{}},1},Array{Float64,1},Array{Set{DynamicPPL.Selector},1}},DynamicPPL.Metadata{Dict{DynamicPPL.VarName{:σ,Tuple{}},Int64},Array{Truncated{Cauchy{Float64},Continuous,Float64},1},Array{DynamicPPL.VarName{:σ,Tuple{}},1},Array{Float64,1},Array{Set{DynamicPPL.Selector},1}},DynamicPPL.Metadata{Dict{DynamicPPL.VarName{:z_Rho,Tuple{}},Int64},Array{DistributionsAD.MatrixOfUnivariate{Continuous,Normal{Float64},FillArrays.Fill{Normal{Float64},2,Tuple{Base.OneTo{Int64},Base.OneTo{Int64}}}},1},Array{DynamicPPL.VarName{:z_Rho,Tuple{}},1},Array{Float64,1},Array{Set{DynamicPPL.Selector},1}}}},Float64}; init_params::Nothing, nadapts::Int64, kwargs::Base.Iterators.Pairs{Union{},Union{},Tuple{},NamedTuple{(),Tuple{}}}) at /Users/stephenjwild/.julia/packages/Turing/XLLTf/src/inference/hmc.jl:167
[32] step(::Random._GLOBAL_RNG, ::DynamicPPL.Model{var"#25#26",(:y, :X, :Z, :gr_id, :n_gr, :m_z),(),(),Tuple{Array{Float64,1},Array{Any,1},Array{Any,1},Array{Int64,1},Int64,Int64},Tuple{}}, ::DynamicPPL.Sampler{NUTS{Turing.Core.ZygoteAD,(),AdvancedHMC.DiagEuclideanMetric}}; resume_from::Nothing, kwargs::Base.Iterators.Pairs{Symbol,Int64,Tuple{Symbol},NamedTuple{(:nadapts,),Tuple{Int64}}}) at /Users/stephenjwild/.julia/packages/DynamicPPL/wf0dU/src/sampler.jl:83
[33] macro expansion at /Users/stephenjwild/.julia/packages/AbstractMCMC/Nw3Wn/src/sample.jl:78 [inlined]
[34] macro expansion at /Users/stephenjwild/.julia/packages/ProgressLogging/6KXlp/src/ProgressLogging.jl:328 [inlined]
[35] macro expansion at /Users/stephenjwild/.julia/packages/AbstractMCMC/Nw3Wn/src/logging.jl:8 [inlined]
[36] mcmcsample(::Random._GLOBAL_RNG, ::DynamicPPL.Model{var"#25#26",(:y, :X, :Z, :gr_id, :n_gr, :m_z),(),(),Tuple{Array{Float64,1},Array{Any,1},Array{Any,1},Array{Int64,1},Int64,Int64},Tuple{}}, ::DynamicPPL.Sampler{NUTS{Turing.Core.ZygoteAD,(),AdvancedHMC.DiagEuclideanMetric}}, ::Int64; progress::Bool, progressname::String, callback::Nothing, discard_initial::Int64, thinning::Int64, chain_type::Type{T} where T, kwargs::Base.Iterators.Pairs{Symbol,Int64,Tuple{Symbol},NamedTuple{(:nadapts,),Tuple{Int64}}}) at /Users/stephenjwild/.julia/packages/AbstractMCMC/Nw3Wn/src/sample.jl:76
[37] sample(::Random._GLOBAL_RNG, ::DynamicPPL.Model{var"#25#26",(:y, :X, :Z, :gr_id, :n_gr, :m_z),(),(),Tuple{Array{Float64,1},Array{Any,1},Array{Any,1},Array{Int64,1},Int64,Int64},Tuple{}}, ::DynamicPPL.Sampler{NUTS{Turing.Core.ZygoteAD,(),AdvancedHMC.DiagEuclideanMetric}}, ::Int64; chain_type::Type{T} where T, resume_from::Nothing, progress::Bool, nadapts::Int64, discard_adapt::Bool, discard_initial::Int64, kwargs::Base.Iterators.Pairs{Union{},Union{},Tuple{},NamedTuple{(),Tuple{}}}) at /Users/stephenjwild/.julia/packages/Turing/XLLTf/src/inference/hmc.jl:133
[38] sample at /Users/stephenjwild/.julia/packages/Turing/XLLTf/src/inference/hmc.jl:116 [inlined]
[39] #sample#2 at /Users/stephenjwild/.julia/packages/Turing/XLLTf/src/inference/Inference.jl:141 [inlined]
[40] sample at /Users/stephenjwild/.julia/packages/Turing/XLLTf/src/inference/Inference.jl:141 [inlined]
[41] #sample#1 at /Users/stephenjwild/.julia/packages/Turing/XLLTf/src/inference/Inference.jl:131 [inlined]
[42] sample(::DynamicPPL.Model{var"#25#26",(:y, :X, :Z, :gr_id, :n_gr, :m_z),(),(),Tuple{Array{Float64,1},Array{Any,1},Array{Any,1},Array{Int64,1},Int64,Int64},Tuple{}}, ::NUTS{Turing.Core.ZygoteAD,(),AdvancedHMC.DiagEuclideanMetric}, ::Int64) at /Users/stephenjwild/.julia/packages/Turing/XLLTf/src/inference/Inference.jl:131
[43] top-level scope at ./timing.jl:174 [inlined]
[44] top-level scope at ./REPL[61]:0

Anything I can do to fix that? So far as I can tell, issues with NamedTuple and Zygote have happened before (see here, for example), but I'm not sure if there's anything I can do from my end with Turing.

Apologies if that seems like a newbie issue. I am relatively new at using Julia and still trying to get my head around Julia and Turing. The error messages are generally more interpretable than R, but it's like reading a foreign language :-p

[devmotion]

I assume there's not much I can do at my end to keep the LKJ prior and still use Tracker as the backend, correct (at least with the way the model is currently coded)?

No, it seems it requires a fix in Bijectors.

Anything I can do to fix that?

It seems the problem is diagm (or rather its gradient definitions in Zygote). Maybe you can avoid the issues by using Diagonal instead of diagm - in contrast to diagm it is a lazy diagonal matrix that does not allocate a full matrix.

[sjwild]

Thanks. Using Diagonal instead of diagm does seem to make a difference. I tired it Diagonal a few times, but in at least one run (two, I think) I got an error message that L_Ρ was not positive definite. Anyways, I am able to use Zygote, but, at least with this particular dataset, is extremely slow compared to forwarddiff and Tracker.

In case someone else has a similar problem and stumbles across this issue, as a second workaround, you can also build your own diagonal matrix, D. I tried it and it works with Zygote as well. Instead of

μ = LinearAlgebra.diagm(τ) * L_Ρ * z_Rho

you can use

D = I(m_z) .* LinearAlgebra.kron(ones(m_z)', τ)
μ = D * L_Ρ * z_Rho

So what you get is

@model function varying_slopes(y, X, Z, gr_id, n_gr, m_z)

    # priors
    Ρ ~ LKJ(m_z, 2.)
    τ ~ filldist(truncated(Cauchy(0, 2), 0, Inf), m_z)
    β ~ Normal(0, 5)
    α ~ Normal(0, 5)
    σ ~ truncated(Cauchy(0, 2), 0, Inf)
    z_Rho ~ filldist(Normal(0, 1), m_z, n_gr)

    # build Rho and random effects
    Ρ = (Ρ' + Ρ) / 2
    L_Ρ = LinearAlgebra.cholesky(Ρ).L
    D = I(m_z) .* LinearAlgebra.kron(ones(m_z)', τ)
    μ = D * L_Ρ * z_Rho

    μ₁ = μ[1, :]
    μ₂ = μ[2, :]

    # model
    ŷ = α .+ μ₁[gr_id] .+ X .* β .+ Z .* μ₂[gr_id]
    y ~ MvNormal(ŷ, σ)

end

It works, albeit a bit slower than using Diagonal or diagm.

Thanks for the help. I appreciate it.

[sjwild]

I'm back!

I have a question for you (it may belong here, maybe under Zygote. I'm not sure). I've tried to run parallel chains using MCMCDistributed. Unfortunately, when I try, I tend to get one chain out four that says

ERROR: PosDefException: matrix is not positive definite; Cholesky factorization failed.

I'm wondering if there is anything I can do to address this? I'd like to be able to run parallel chains.

[devmotion]

Is the error consistently reproducible with parallel sampling but never occurs with standard sampling? In principle, this numerical error should not be related to parallel sampling and be reproducible also with regular sampling. Parallel sampling should completely decouple the model and samplers of different processes (MCMCDistributed) or threads (MCMCThreads), and each chain is sampled independently with a randomly chosen seed.

[sjwild]

Good question! I should have been clearer in my post. When I run a single chain using Zygote as the backend, I still get the same error message once in a while (maybe 10% - 20% of the time. I haven't kept track, to be honest). When running parallel chains, I have yet to run an instance without the error message.

I haven't had an issue with the PosDefException using forwarddiff.

[devmotion]

I don't know why you would get numerical errors with Zygote but not with ForwardDiff - this seems weird.

[sjwild]

I agree. I initially thought it might be related to this issue and solution, but it seemed weird it only happened for me with Zygote and not with forwarddiff. Hence why I thought I should ask.

If it is the case that it's related to the linked issue, then setting reasonable initial values is probably the solution.

[storopoli]

I am also having problems with Zygote in this model. Turing with the standard backend is OK.

sing Turing, DataFrames, Pipe, CSV, HTTP, Zygote, ReverseDiff
using Random:seed!
using Statistics: mean, std

seed!(1)
setprogress!(true)

df = @pipe HTTP.get("https://github.com/selva86/datasets/blob/master/mpg_ggplot2.csv?raw=TRUE").body |>
    CSV.read(_, DataFrame)

#### Data Prep ####
idx_map = Dict(key => idx for (idx, key) in enumerate(unique(df.class)))
y = df[:, :hwy]
idx = getindex.(Ref(idx_map), df.class)
X = Matrix(select(df, [:displ, :year])) # the model matrix

### NCP Varying Intercept Model ####
@model varying_intercept_ncp(X, idx, y; n_gr=length(unique(idx)), predictors=size(X, 2)) = begin
    # priors
    μ ~ Normal(mean(y), 2.5 * std(y))       # population-level intercept
    σ ~ Exponential(1 / std(y))             # residual SD
    # Coefficients Student-t(ν = 3)
    β ~ filldist(TDist(3), predictors)
    # Prior for variance of random intercepts. Usually requires thoughtful specification.
    σⱼ ~ Truncated(Cauchy(0, 2), 0, Inf)
    zⱼ ~ filldist(Normal(0, 1), n_gr)      # NCP group-level intercepts

    # likelihood
    ŷ = μ .+ X * β .+ zⱼ[idx] .* σⱼ
    y ~ MvNormal(ŷ, σ)
end

model_ncp = varying_intercept_ncp(X, idx, y)
Turing.setadbackend(:zygote)
chn_zygote = sample(model_ncp, NUTS(1_000, 0.65), 2_000)

I get this error:

ERROR: LoadError: Compiling Tuple{Base.var"##depwarn#868", Bool, typeof(Base.depwarn), String, Symbol}: try/catch is not supported.
Stacktrace:
  [1] error(s::String)
    @ Base ./error.jl:33
  [2] instrument(ir::IRTools.Inner.IR)
    @ Zygote ~/.julia/packages/Zygote/KpME9/src/compiler/reverse.jl:121
  [3] #Primal#20
    @ ~/.julia/packages/Zygote/KpME9/src/compiler/reverse.jl:202 [inlined]
  [4] Zygote.Adjoint(ir::IRTools.Inner.IR; varargs::Nothing, normalise::Bool)
    @ Zygote ~/.julia/packages/Zygote/KpME9/src/compiler/reverse.jl:315
  [5] _lookup_grad(T::Type)
    @ Zygote ~/.julia/packages/Zygote/KpME9/src/compiler/emit.jl:101
  [6] #s2930#1123
    @ ~/.julia/packages/Zygote/KpME9/src/compiler/interface2.jl:34 [inlined]
  [7] var"#s2930#1123"(T::Any, j::Any, Δ::Any)
    @ Zygote ./none:0
  [8] (::Core.GeneratedFunctionStub)(::Any, ::Vararg{Any, N} where N)
    @ Core ./boot.jl:571
  [9] Pullback
    @ ./deprecated.jl:80 [inlined]
 [10] (::typeof(∂(depwarn)))(Δ::Nothing)
    @ Zygote ~/.julia/packages/Zygote/KpME9/src/compiler/interface2.jl:0
 [11] Pullback
    @ ./deprecated.jl:71 [inlined]
 [12] (::typeof(∂(Truncated)))(Δ::NamedTuple{(:untruncated, :lower, :upper, :lcdf, :ucdf, :tp, :logtp), Tuple{NamedTuple{(:μ, :σ), Tuple{Float64, Float64}}, Float64, Float64, Nothing, Nothing, Nothing, Int64}})
    @ Zygote ~/.julia/packages/Zygote/KpME9/src/compiler/interface2.jl:0
 [13] Pullback
    @ ~/Documents/Julia_Scripts/Turing/Turing_hierarchical.jl:51 [inlined]
 [14] (::typeof(∂(#8)))(Δ::Nothing)
    @ Zygote ~/.julia/packages/Zygote/KpME9/src/compiler/interface2.jl:0
 [15] macro expansion
    @ ~/.julia/packages/DynamicPPL/wf0dU/src/model.jl:0 [inlined]
 [16] Pullback
    @ ~/.julia/packages/DynamicPPL/wf0dU/src/model.jl:154 [inlined]
 [17] (::typeof(∂(_evaluate)))(Δ::Nothing)
    @ Zygote ~/.julia/packages/Zygote/KpME9/src/compiler/interface2.jl:0
 [18] Pullback
    @ ~/.julia/packages/DynamicPPL/wf0dU/src/model.jl:144 [inlined]
 [19] (::typeof(∂(evaluate_threadsafe)))(Δ::Nothing)
    @ Zygote ~/.julia/packages/Zygote/KpME9/src/compiler/interface2.jl:0
 [20] Pullback
    @ ~/.julia/packages/DynamicPPL/wf0dU/src/model.jl:94 [inlined]
 [21] (::typeof(∂(λ)))(Δ::Nothing)
    @ Zygote ~/.julia/packages/Zygote/KpME9/src/compiler/interface2.jl:0
 [22] #151
    @ ~/.julia/packages/Zygote/KpME9/src/lib/lib.jl:191 [inlined]
 [23] #1682#back
    @ ~/.julia/packages/ZygoteRules/OjfTt/src/adjoint.jl:59 [inlined]
 [24] Pullback
    @ ~/.julia/packages/DynamicPPL/wf0dU/src/model.jl:98 [inlined]
 [25] Pullback
    @ ~/.julia/packages/Turing/uAz5c/src/core/ad.jl:165 [inlined]
 [26] (::typeof(∂(λ)))(Δ::Int64)
    @ Zygote ~/.julia/packages/Zygote/KpME9/src/compiler/interface2.jl:0
 [27] (::Zygote.var"#41#42"{typeof(∂(λ))})(Δ::Int64)
    @ Zygote ~/.julia/packages/Zygote/KpME9/src/compiler/interface.jl:40
 [28] gradient_logp(backend::Turing.Core.ZygoteAD, θ::Vector{Float64}, vi::DynamicPPL.TypedVarInfo{NamedTuple{(:μ, :σ, :β, :σⱼ, :zⱼ), Tuple{DynamicPPL.Metadata{Dict{DynamicPPL.VarName{:μ, Tuple{}}, Int64}, Vector{Normal{Float64}}, Vector{DynamicPPL.VarName{:μ, Tuple{}}}, Vector{Float64}, Vector{Set{DynamicPPL.Selector}}}, DynamicPPL.Metadata{Dict{DynamicPPL.VarName{:σ, Tuple{}}, Int64}, Vector{Exponential{Float64}}, Vector{DynamicPPL.VarName{:σ, Tuple{}}}, Vector{Float64}, Vector{Set{DynamicPPL.Selector}}}, DynamicPPL.Metadata{Dict{DynamicPPL.VarName{:β, Tuple{}}, Int64}, Vector{Product{Continuous, TDist{Float64}, FillArrays.Fill{TDist{Float64}, 1, Tuple{Base.OneTo{Int64}}}}}, Vector{DynamicPPL.VarName{:β, Tuple{}}}, Vector{Float64}, Vector{Set{DynamicPPL.Selector}}}, DynamicPPL.Metadata{Dict{DynamicPPL.VarName{:σⱼ, Tuple{}}, Int64}, Vector{Truncated{Cauchy{Float64}, Continuous, Float64}}, Vector{DynamicPPL.VarName{:σⱼ, Tuple{}}}, Vector{Float64}, Vector{Set{DynamicPPL.Selector}}}, DynamicPPL.Metadata{Dict{DynamicPPL.VarName{:zⱼ, Tuple{}}, Int64}, Vector{DistributionsAD.TuringScalMvNormal{Vector{Float64}, Float64}}, Vector{DynamicPPL.VarName{:zⱼ, Tuple{}}}, Vector{Float64}, Vector{Set{DynamicPPL.Selector}}}}}, Float64}, model::DynamicPPL.Model{var"#8#9", (:X, :idx, :y, :n_gr, :predictors), (:n_gr, :predictors), (), Tuple{Matrix{Float64}, Vector{Int64}, Vector{Int64}, Int64, Int64}, Tuple{Int64, Int64}}, sampler::DynamicPPL.Sampler{NUTS{Turing.Core.ZygoteAD, (), AdvancedHMC.DiagEuclideanMetric}}, context::DynamicPPL.DefaultContext)
    @ Turing.Core ~/.julia/packages/Turing/uAz5c/src/core/ad.jl:171
 [29] gradient_logp
    @ ~/.julia/packages/Turing/uAz5c/src/core/ad.jl:83 [inlined]
 [30] ∂logπ∂θ
    @ ~/.julia/packages/Turing/uAz5c/src/inference/hmc.jl:482 [inlined]
 [31] ∂H∂θ
    @ ~/.julia/packages/AdvancedHMC/MIxdK/src/hamiltonian.jl:31 [inlined]
 [32] phasepoint
    @ ~/.julia/packages/AdvancedHMC/MIxdK/src/hamiltonian.jl:69 [inlined]
 [33] phasepoint(rng::Random._GLOBAL_RNG, θ::Vector{Float64}, h::AdvancedHMC.Hamiltonian{AdvancedHMC.DiagEuclideanMetric{Float64, Vector{Float64}}, Turing.Inference.var"#logπ#52"{DynamicPPL.TypedVarInfo{NamedTuple{(:μ, :σ, :β, :σⱼ, :zⱼ), Tuple{DynamicPPL.Metadata{Dict{DynamicPPL.VarName{:μ, Tuple{}}, Int64}, Vector{Normal{Float64}}, Vector{DynamicPPL.VarName{:μ, Tuple{}}}, Vector{Float64}, Vector{Set{DynamicPPL.Selector}}}, DynamicPPL.Metadata{Dict{DynamicPPL.VarName{:σ, Tuple{}}, Int64}, Vector{Exponential{Float64}}, Vector{DynamicPPL.VarName{:σ, Tuple{}}}, Vector{Float64}, Vector{Set{DynamicPPL.Selector}}}, DynamicPPL.Metadata{Dict{DynamicPPL.VarName{:β, Tuple{}}, Int64}, Vector{Product{Continuous, TDist{Float64}, FillArrays.Fill{TDist{Float64}, 1, Tuple{Base.OneTo{Int64}}}}}, Vector{DynamicPPL.VarName{:β, Tuple{}}}, Vector{Float64}, Vector{Set{DynamicPPL.Selector}}}, DynamicPPL.Metadata{Dict{DynamicPPL.VarName{:σⱼ, Tuple{}}, Int64}, Vector{Truncated{Cauchy{Float64}, Continuous, Float64}}, Vector{DynamicPPL.VarName{:σⱼ, Tuple{}}}, Vector{Float64}, Vector{Set{DynamicPPL.Selector}}}, DynamicPPL.Metadata{Dict{DynamicPPL.VarName{:zⱼ, Tuple{}}, Int64}, Vector{DistributionsAD.TuringScalMvNormal{Vector{Float64}, Float64}}, Vector{DynamicPPL.VarName{:zⱼ, Tuple{}}}, Vector{Float64}, Vector{Set{DynamicPPL.Selector}}}}}, Float64}, DynamicPPL.Sampler{NUTS{Turing.Core.ZygoteAD, (), AdvancedHMC.DiagEuclideanMetric}}, DynamicPPL.Model{var"#8#9", (:X, :idx, :y, :n_gr, :predictors), (:n_gr, :predictors), (), Tuple{Matrix{Float64}, Vector{Int64}, Vector{Int64}, Int64, Int64}, Tuple{Int64, Int64}}}, Turing.Inference.var"#∂logπ∂θ#51"{DynamicPPL.TypedVarInfo{NamedTuple{(:μ, :σ, :β, :σⱼ, :zⱼ), Tuple{DynamicPPL.Metadata{Dict{DynamicPPL.VarName{:μ, Tuple{}}, Int64}, Vector{Normal{Float64}}, Vector{DynamicPPL.VarName{:μ, Tuple{}}}, Vector{Float64}, Vector{Set{DynamicPPL.Selector}}}, DynamicPPL.Metadata{Dict{DynamicPPL.VarName{:σ, Tuple{}}, Int64}, Vector{Exponential{Float64}}, Vector{DynamicPPL.VarName{:σ, Tuple{}}}, Vector{Float64}, Vector{Set{DynamicPPL.Selector}}}, DynamicPPL.Metadata{Dict{DynamicPPL.VarName{:β, Tuple{}}, Int64}, Vector{Product{Continuous, TDist{Float64}, FillArrays.Fill{TDist{Float64}, 1, Tuple{Base.OneTo{Int64}}}}}, Vector{DynamicPPL.VarName{:β, Tuple{}}}, Vector{Float64}, Vector{Set{DynamicPPL.Selector}}}, DynamicPPL.Metadata{Dict{DynamicPPL.VarName{:σⱼ, Tuple{}}, Int64}, Vector{Truncated{Cauchy{Float64}, Continuous, Float64}}, Vector{DynamicPPL.VarName{:σⱼ, Tuple{}}}, Vector{Float64}, Vector{Set{DynamicPPL.Selector}}}, DynamicPPL.Metadata{Dict{DynamicPPL.VarName{:zⱼ, Tuple{}}, Int64}, Vector{DistributionsAD.TuringScalMvNormal{Vector{Float64}, Float64}}, Vector{DynamicPPL.VarName{:zⱼ, Tuple{}}}, Vector{Float64}, Vector{Set{DynamicPPL.Selector}}}}}, Float64}, DynamicPPL.Sampler{NUTS{Turing.Core.ZygoteAD, (), AdvancedHMC.DiagEuclideanMetric}}, DynamicPPL.Model{var"#8#9", (:X, :idx, :y, :n_gr, :predictors), (:n_gr, :predictors), (), Tuple{Matrix{Float64}, Vector{Int64}, Vector{Int64}, Int64, Int64}, Tuple{Int64, Int64}}}})
    @ AdvancedHMC ~/.julia/packages/AdvancedHMC/MIxdK/src/hamiltonian.jl:139
 [34] initialstep(rng::Random._GLOBAL_RNG, model::DynamicPPL.Model{var"#8#9", (:X, :idx, :y, :n_gr, :predictors), (:n_gr, :predictors), (), Tuple{Matrix{Float64}, Vector{Int64}, Vector{Int64}, Int64, Int64}, Tuple{Int64, Int64}}, spl::DynamicPPL.Sampler{NUTS{Turing.Core.ZygoteAD, (), AdvancedHMC.DiagEuclideanMetric}}, vi::DynamicPPL.TypedVarInfo{NamedTuple{(:μ, :σ, :β, :σⱼ, :zⱼ), Tuple{DynamicPPL.Metadata{Dict{DynamicPPL.VarName{:μ, Tuple{}}, Int64}, Vector{Normal{Float64}}, Vector{DynamicPPL.VarName{:μ, Tuple{}}}, Vector{Float64}, Vector{Set{DynamicPPL.Selector}}}, DynamicPPL.Metadata{Dict{DynamicPPL.VarName{:σ, Tuple{}}, Int64}, Vector{Exponential{Float64}}, Vector{DynamicPPL.VarName{:σ, Tuple{}}}, Vector{Float64}, Vector{Set{DynamicPPL.Selector}}}, DynamicPPL.Metadata{Dict{DynamicPPL.VarName{:β, Tuple{}}, Int64}, Vector{Product{Continuous, TDist{Float64}, FillArrays.Fill{TDist{Float64}, 1, Tuple{Base.OneTo{Int64}}}}}, Vector{DynamicPPL.VarName{:β, Tuple{}}}, Vector{Float64}, Vector{Set{DynamicPPL.Selector}}}, DynamicPPL.Metadata{Dict{DynamicPPL.VarName{:σⱼ, Tuple{}}, Int64}, Vector{Truncated{Cauchy{Float64}, Continuous, Float64}}, Vector{DynamicPPL.VarName{:σⱼ, Tuple{}}}, Vector{Float64}, Vector{Set{DynamicPPL.Selector}}}, DynamicPPL.Metadata{Dict{DynamicPPL.VarName{:zⱼ, Tuple{}}, Int64}, Vector{DistributionsAD.TuringScalMvNormal{Vector{Float64}, Float64}}, Vector{DynamicPPL.VarName{:zⱼ, Tuple{}}}, Vector{Float64}, Vector{Set{DynamicPPL.Selector}}}}}, Float64}; init_params::Nothing, nadapts::Int64, kwargs::Base.Iterators.Pairs{Union{}, Union{}, Tuple{}, NamedTuple{(), Tuple{}}})
    @ Turing.Inference ~/.julia/packages/Turing/uAz5c/src/inference/hmc.jl:174
 [35] step(rng::Random._GLOBAL_RNG, model::DynamicPPL.Model{var"#8#9", (:X, :idx, :y, :n_gr, :predictors), (:n_gr, :predictors), (), Tuple{Matrix{Float64}, Vector{Int64}, Vector{Int64}, Int64, Int64}, Tuple{Int64, Int64}}, spl::DynamicPPL.Sampler{NUTS{Turing.Core.ZygoteAD, (), AdvancedHMC.DiagEuclideanMetric}}; resume_from::Nothing, kwargs::Base.Iterators.Pairs{Symbol, Int64, Tuple{Symbol}, NamedTuple{(:nadapts,), Tuple{Int64}}})
    @ DynamicPPL ~/.julia/packages/DynamicPPL/wf0dU/src/sampler.jl:83
 [36] macro expansion
    @ ~/.julia/packages/AbstractMCMC/Nw3Wn/src/sample.jl:78 [inlined]
 [37] macro expansion
    @ ~/.julia/packages/ProgressLogging/6KXlp/src/ProgressLogging.jl:328 [inlined]
 [38] macro expansion
    @ ~/.julia/packages/AbstractMCMC/Nw3Wn/src/logging.jl:8 [inlined]
 [39] mcmcsample(rng::Random._GLOBAL_RNG, model::DynamicPPL.Model{var"#8#9", (:X, :idx, :y, :n_gr, :predictors), (:n_gr, :predictors), (), Tuple{Matrix{Float64}, Vector{Int64}, Vector{Int64}, Int64, Int64}, Tuple{Int64, Int64}}, 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.Iterators.Pairs{Symbol, Int64, Tuple{Symbol}, NamedTuple{(:nadapts,), Tuple{Int64}}})
    @ AbstractMCMC ~/.julia/packages/AbstractMCMC/Nw3Wn/src/sample.jl:76
 [40] sample(rng::Random._GLOBAL_RNG, model::DynamicPPL.Model{var"#8#9", (:X, :idx, :y, :n_gr, :predictors), (:n_gr, :predictors), (), Tuple{Matrix{Float64}, Vector{Int64}, Vector{Int64}, Int64, Int64}, Tuple{Int64, Int64}}, 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.Iterators.Pairs{Union{}, Union{}, Tuple{}, NamedTuple{(), Tuple{}}})
    @ Turing.Inference ~/.julia/packages/Turing/uAz5c/src/inference/hmc.jl:140
 [41] sample
    @ ~/.julia/packages/Turing/uAz5c/src/inference/hmc.jl:123 [inlined]
 [42] #sample#2
    @ ~/.julia/packages/Turing/uAz5c/src/inference/Inference.jl:142 [inlined]
 [43] sample
    @ ~/.julia/packages/Turing/uAz5c/src/inference/Inference.jl:142 [inlined]
 [44] #sample#1
    @ ~/.julia/packages/Turing/uAz5c/src/inference/Inference.jl:132 [inlined]
 [45] sample(model::DynamicPPL.Model{var"#8#9", (:X, :idx, :y, :n_gr, :predictors), (:n_gr, :predictors), (), Tuple{Matrix{Float64}, Vector{Int64}, Vector{Int64}, Int64, Int64}, Tuple{Int64, Int64}}, alg::NUTS{Turing.Core.ZygoteAD, (), AdvancedHMC.DiagEuclideanMetric}, N::Int64)
    @ Turing.Inference ~/.julia/packages/Turing/uAz5c/src/inference/Inference.jl:132
 [46] top-level scope
    @ ./timing.jl:206 [inlined]
 [47] top-level scope
    @ ~/Documents/Julia_Scripts/Turing/Turing_hierarchical.jl:0
in expression starting at /Users/storopoli/Documents/Julia_Scripts/Turing/Turing_hierarchical.jl:68

Seems related with the Truncated Cauchy in the model:

σⱼ ~ Truncated(Cauchy(0, 2), 0, Inf)

[fce5fe82] Turing v0.15.1
[e88e6eb3] Zygote v0.6.3

[devmotion]

Zygote does not allow to differentiate through try-catch-blocks, and therefore also not if there are logging statements such as @warn (they use a try-catch-block internally: https://github.com/JuliaLang/julia/blob/c79309bffa2842f7864d06fc2b135e79da1daec1/base/logging.jl#L339-L346). In your example, the logging statements are part of the deprecation warnings (https://github.com/JuliaLang/julia/blob/c79309bffa2842f7864d06fc2b135e79da1daec1/base/deprecated.jl#L204) which are caused by the use of the deprecated constructor Truncated (https://github.com/JuliaStats/Distributions.jl/blob/7120d503c269503b0037a047399452b596ab43d4/src/truncate.jl#L72).

As a general rule, one should use truncated instead of Truncated: Truncated always creates instances of Truncated whereas truncated can create possibly more efficient objects (such as TruncatedNormal etc.).

[storopoli]

Changed to

σⱼ ~ truncated(Cauchy(0, 2), 0, Inf)

But still getting:

julia> @time chn_zygote = sample(model_ncp, NUTS(1_000, 0.65), MCMCThreads(), 2_000, 4)
ERROR: LoadError: TaskFailedException

    nested task error: TaskFailedException
    Stacktrace:
     [1] wait
       @ ./task.jl:317 [inlined]
     [2] threading_run(func::Function)
       @ Base.Threads ./threadingconstructs.jl:34
     [3] macro expansion
       @ ./threadingconstructs.jl:93 [inlined]
     [4] macro expansion
       @ ~/.julia/packages/AbstractMCMC/Nw3Wn/src/sample.jl:271 [inlined]
     [5] (::AbstractMCMC.var"#31#44"{Bool, Base.Iterators.Pairs{Symbol, UnionAll, Tuple{Symbol}, NamedTuple{(:chain_type,), Tuple{UnionAll}}}, Int64, Int64, Vector{Any}, Vector{UInt64}, Vector{DynamicPPL.Sampler{NUTS{Turing.Core.ZygoteAD, (), AdvancedHMC.DiagEuclideanMetric}}}, Vector{DynamicPPL.Model{var"#17#18", (:X, :idx, :y, :n_gr, :predictors), (:n_gr, :predictors), (), Tuple{Matrix{Float64}, Vector{Int64}, Vector{Int64}, Int64, Int64}, Tuple{Int64, Int64}}}, Vector{Random._GLOBAL_RNG}})()
       @ AbstractMCMC ./task.jl:406
    
        nested task error: InexactError: Int64(0.008583690987124463)
        Stacktrace:
          [1] Int64
            @ ./float.jl:723 [inlined]
          [2] convert
            @ ./number.jl:7 [inlined]
          [3] _backvar(xs::Vector{Int64}, Δ::Float64, N::Int64, mean::Float64)
            @ Zygote ~/.julia/packages/Zygote/KpME9/src/lib/array.jl:315
          [4] _backvar
            @ ~/.julia/packages/Zygote/KpME9/src/lib/array.jl:314 [inlined]
          [5] #627
            @ ~/.julia/packages/Zygote/KpME9/src/lib/array.jl:319 [inlined]
          [6] (::Zygote.var"#2763#back#629"{Zygote.var"#627#628"{Bool, Colon, Float64, Vector{Int64}, Float64}})(Δ::Float64)
            @ Zygote ~/.julia/packages/ZygoteRules/OjfTt/src/adjoint.jl:59
          [7] Pullback
            @ ~/Documents/Julia_Scripts/Turing/Turing_hierarchical.jl:47 [inlined]
          [8] (::typeof(∂(#17)))(Δ::Nothing)
            @ Zygote ~/.julia/packages/Zygote/KpME9/src/compiler/interface2.jl:0
          [9] macro expansion
            @ ~/.julia/packages/DynamicPPL/wf0dU/src/model.jl:0 [inlined]
         [10] Pullback
            @ ~/.julia/packages/DynamicPPL/wf0dU/src/model.jl:154 [inlined]
         [11] (::typeof(∂(_evaluate)))(Δ::Nothing)
            @ Zygote ~/.julia/packages/Zygote/KpME9/src/compiler/interface2.jl:0
         [12] Pullback
            @ ~/.julia/packages/DynamicPPL/wf0dU/src/model.jl:144 [inlined]
         [13] (::typeof(∂(evaluate_threadsafe)))(Δ::Nothing)
            @ Zygote ~/.julia/packages/Zygote/KpME9/src/compiler/interface2.jl:0
         [14] Pullback
            @ ~/.julia/packages/DynamicPPL/wf0dU/src/model.jl:94 [inlined]
         [15] (::typeof(∂(λ)))(Δ::Nothing)
            @ Zygote ~/.julia/packages/Zygote/KpME9/src/compiler/interface2.jl:0
         [16] #151
            @ ~/.julia/packages/Zygote/KpME9/src/lib/lib.jl:191 [inlined]
         [17] #1682#back
            @ ~/.julia/packages/ZygoteRules/OjfTt/src/adjoint.jl:59 [inlined]
         [18] Pullback
            @ ~/.julia/packages/DynamicPPL/wf0dU/src/model.jl:98 [inlined]
         [19] Pullback
            @ ~/.julia/packages/Turing/uAz5c/src/core/ad.jl:165 [inlined]
         [20] (::typeof(∂(λ)))(Δ::Int64)
            @ Zygote ~/.julia/packages/Zygote/KpME9/src/compiler/interface2.jl:0
         [21] (::Zygote.var"#41#42"{typeof(∂(λ))})(Δ::Int64)
            @ Zygote ~/.julia/packages/Zygote/KpME9/src/compiler/interface.jl:40
         [22] gradient_logp(backend::Turing.Core.ZygoteAD, θ::Vector{Float64}, vi::DynamicPPL.TypedVarInfo{NamedTuple{(:μ, :σ, :β, :σⱼ, :zⱼ), Tuple{DynamicPPL.Metadata{Dict{DynamicPPL.VarName{:μ, Tuple{}}, Int64}, Vector{Normal{Float64}}, Vector{DynamicPPL.VarName{:μ, Tuple{}}}, Vector{Float64}, Vector{Set{DynamicPPL.Selector}}}, DynamicPPL.Metadata{Dict{DynamicPPL.VarName{:σ, Tuple{}}, Int64}, Vector{Exponential{Float64}}, Vector{DynamicPPL.VarName{:σ, Tuple{}}}, Vector{Float64}, Vector{Set{DynamicPPL.Selector}}}, DynamicPPL.Metadata{Dict{DynamicPPL.VarName{:β, Tuple{}}, Int64}, Vector{Product{Continuous, TDist{Float64}, FillArrays.Fill{TDist{Float64}, 1, Tuple{Base.OneTo{Int64}}}}}, Vector{DynamicPPL.VarName{:β, Tuple{}}}, Vector{Float64}, Vector{Set{DynamicPPL.Selector}}}, DynamicPPL.Metadata{Dict{DynamicPPL.VarName{:σⱼ, Tuple{}}, Int64}, Vector{Truncated{Cauchy{Float64}, Continuous, Float64}}, Vector{DynamicPPL.VarName{:σⱼ, Tuple{}}}, Vector{Float64}, Vector{Set{DynamicPPL.Selector}}}, DynamicPPL.Metadata{Dict{DynamicPPL.VarName{:zⱼ, Tuple{}}, Int64}, Vector{DistributionsAD.TuringScalMvNormal{Vector{Float64}, Float64}}, Vector{DynamicPPL.VarName{:zⱼ, Tuple{}}}, Vector{Float64}, Vector{Set{DynamicPPL.Selector}}}}}, Float64}, model::DynamicPPL.Model{var"#17#18", (:X, :idx, :y, :n_gr, :predictors), (:n_gr, :predictors), (), Tuple{Matrix{Float64}, Vector{Int64}, Vector{Int64}, Int64, Int64}, Tuple{Int64, Int64}}, sampler::DynamicPPL.Sampler{NUTS{Turing.Core.ZygoteAD, (), AdvancedHMC.DiagEuclideanMetric}}, context::DynamicPPL.DefaultContext)
            @ Turing.Core ~/.julia/packages/Turing/uAz5c/src/core/ad.jl:171
         [23] gradient_logp
            @ ~/.julia/packages/Turing/uAz5c/src/core/ad.jl:83 [inlined]
         [24] ∂logπ∂θ
            @ ~/.julia/packages/Turing/uAz5c/src/inference/hmc.jl:482 [inlined]
         [25] ∂H∂θ
            @ ~/.julia/packages/AdvancedHMC/MIxdK/src/hamiltonian.jl:31 [inlined]
         [26] phasepoint
            @ ~/.julia/packages/AdvancedHMC/MIxdK/src/hamiltonian.jl:69 [inlined]
         [27] phasepoint(rng::Random._GLOBAL_RNG, θ::Vector{Float64}, h::AdvancedHMC.Hamiltonian{AdvancedHMC.DiagEuclideanMetric{Float64, Vector{Float64}}, Turing.Inference.var"#logπ#52"{DynamicPPL.TypedVarInfo{NamedTuple{(:μ, :σ, :β, :σⱼ, :zⱼ), Tuple{DynamicPPL.Metadata{Dict{DynamicPPL.VarName{:μ, Tuple{}}, Int64}, Vector{Normal{Float64}}, Vector{DynamicPPL.VarName{:μ, Tuple{}}}, Vector{Float64}, Vector{Set{DynamicPPL.Selector}}}, DynamicPPL.Metadata{Dict{DynamicPPL.VarName{:σ, Tuple{}}, Int64}, Vector{Exponential{Float64}}, Vector{DynamicPPL.VarName{:σ, Tuple{}}}, Vector{Float64}, Vector{Set{DynamicPPL.Selector}}}, DynamicPPL.Metadata{Dict{DynamicPPL.VarName{:β, Tuple{}}, Int64}, Vector{Product{Continuous, TDist{Float64}, FillArrays.Fill{TDist{Float64}, 1, Tuple{Base.OneTo{Int64}}}}}, Vector{DynamicPPL.VarName{:β, Tuple{}}}, Vector{Float64}, Vector{Set{DynamicPPL.Selector}}}, DynamicPPL.Metadata{Dict{DynamicPPL.VarName{:σⱼ, Tuple{}}, Int64}, Vector{Truncated{Cauchy{Float64}, Continuous, Float64}}, Vector{DynamicPPL.VarName{:σⱼ, Tuple{}}}, Vector{Float64}, Vector{Set{DynamicPPL.Selector}}}, DynamicPPL.Metadata{Dict{DynamicPPL.VarName{:zⱼ, Tuple{}}, Int64}, Vector{DistributionsAD.TuringScalMvNormal{Vector{Float64}, Float64}}, Vector{DynamicPPL.VarName{:zⱼ, Tuple{}}}, Vector{Float64}, Vector{Set{DynamicPPL.Selector}}}}}, Float64}, DynamicPPL.Sampler{NUTS{Turing.Core.ZygoteAD, (), AdvancedHMC.DiagEuclideanMetric}}, DynamicPPL.Model{var"#17#18", (:X, :idx, :y, :n_gr, :predictors), (:n_gr, :predictors), (), Tuple{Matrix{Float64}, Vector{Int64}, Vector{Int64}, Int64, Int64}, Tuple{Int64, Int64}}}, Turing.Inference.var"#∂logπ∂θ#51"{DynamicPPL.TypedVarInfo{NamedTuple{(:μ, :σ, :β, :σⱼ, :zⱼ), Tuple{DynamicPPL.Metadata{Dict{DynamicPPL.VarName{:μ, Tuple{}}, Int64}, Vector{Normal{Float64}}, Vector{DynamicPPL.VarName{:μ, Tuple{}}}, Vector{Float64}, Vector{Set{DynamicPPL.Selector}}}, DynamicPPL.Metadata{Dict{DynamicPPL.VarName{:σ, Tuple{}}, Int64}, Vector{Exponential{Float64}}, Vector{DynamicPPL.VarName{:σ, Tuple{}}}, Vector{Float64}, Vector{Set{DynamicPPL.Selector}}}, DynamicPPL.Metadata{Dict{DynamicPPL.VarName{:β, Tuple{}}, Int64}, Vector{Product{Continuous, TDist{Float64}, FillArrays.Fill{TDist{Float64}, 1, Tuple{Base.OneTo{Int64}}}}}, Vector{DynamicPPL.VarName{:β, Tuple{}}}, Vector{Float64}, Vector{Set{DynamicPPL.Selector}}}, DynamicPPL.Metadata{Dict{DynamicPPL.VarName{:σⱼ, Tuple{}}, Int64}, Vector{Truncated{Cauchy{Float64}, Continuous, Float64}}, Vector{DynamicPPL.VarName{:σⱼ, Tuple{}}}, Vector{Float64}, Vector{Set{DynamicPPL.Selector}}}, DynamicPPL.Metadata{Dict{DynamicPPL.VarName{:zⱼ, Tuple{}}, Int64}, Vector{DistributionsAD.TuringScalMvNormal{Vector{Float64}, Float64}}, Vector{DynamicPPL.VarName{:zⱼ, Tuple{}}}, Vector{Float64}, Vector{Set{DynamicPPL.Selector}}}}}, Float64}, DynamicPPL.Sampler{NUTS{Turing.Core.ZygoteAD, (), AdvancedHMC.DiagEuclideanMetric}}, DynamicPPL.Model{var"#17#18", (:X, :idx, :y, :n_gr, :predictors), (:n_gr, :predictors), (), Tuple{Matrix{Float64}, Vector{Int64}, Vector{Int64}, Int64, Int64}, Tuple{Int64, Int64}}}})
            @ AdvancedHMC ~/.julia/packages/AdvancedHMC/MIxdK/src/hamiltonian.jl:139
         [28] initialstep(rng::Random._GLOBAL_RNG, model::DynamicPPL.Model{var"#17#18", (:X, :idx, :y, :n_gr, :predictors), (:n_gr, :predictors), (), Tuple{Matrix{Float64}, Vector{Int64}, Vector{Int64}, Int64, Int64}, Tuple{Int64, Int64}}, spl::DynamicPPL.Sampler{NUTS{Turing.Core.ZygoteAD, (), AdvancedHMC.DiagEuclideanMetric}}, vi::DynamicPPL.TypedVarInfo{NamedTuple{(:μ, :σ, :β, :σⱼ, :zⱼ), Tuple{DynamicPPL.Metadata{Dict{DynamicPPL.VarName{:μ, Tuple{}}, Int64}, Vector{Normal{Float64}}, Vector{DynamicPPL.VarName{:μ, Tuple{}}}, Vector{Float64}, Vector{Set{DynamicPPL.Selector}}}, DynamicPPL.Metadata{Dict{DynamicPPL.VarName{:σ, Tuple{}}, Int64}, Vector{Exponential{Float64}}, Vector{DynamicPPL.VarName{:σ, Tuple{}}}, Vector{Float64}, Vector{Set{DynamicPPL.Selector}}}, DynamicPPL.Metadata{Dict{DynamicPPL.VarName{:β, Tuple{}}, Int64}, Vector{Product{Continuous, TDist{Float64}, FillArrays.Fill{TDist{Float64}, 1, Tuple{Base.OneTo{Int64}}}}}, Vector{DynamicPPL.VarName{:β, Tuple{}}}, Vector{Float64}, Vector{Set{DynamicPPL.Selector}}}, DynamicPPL.Metadata{Dict{DynamicPPL.VarName{:σⱼ, Tuple{}}, Int64}, Vector{Truncated{Cauchy{Float64}, Continuous, Float64}}, Vector{DynamicPPL.VarName{:σⱼ, Tuple{}}}, Vector{Float64}, Vector{Set{DynamicPPL.Selector}}}, DynamicPPL.Metadata{Dict{DynamicPPL.VarName{:zⱼ, Tuple{}}, Int64}, Vector{DistributionsAD.TuringScalMvNormal{Vector{Float64}, Float64}}, Vector{DynamicPPL.VarName{:zⱼ, Tuple{}}}, Vector{Float64}, Vector{Set{DynamicPPL.Selector}}}}}, Float64}; init_params::Nothing, nadapts::Int64, kwargs::Base.Iterators.Pairs{Union{}, Union{}, Tuple{}, NamedTuple{(), Tuple{}}})
            @ Turing.Inference ~/.julia/packages/Turing/uAz5c/src/inference/hmc.jl:174
         [29] step(rng::Random._GLOBAL_RNG, model::DynamicPPL.Model{var"#17#18", (:X, :idx, :y, :n_gr, :predictors), (:n_gr, :predictors), (), Tuple{Matrix{Float64}, Vector{Int64}, Vector{Int64}, Int64, Int64}, Tuple{Int64, Int64}}, spl::DynamicPPL.Sampler{NUTS{Turing.Core.ZygoteAD, (), AdvancedHMC.DiagEuclideanMetric}}; resume_from::Nothing, kwargs::Base.Iterators.Pairs{Symbol, Int64, Tuple{Symbol}, NamedTuple{(:nadapts,), Tuple{Int64}}})
            @ DynamicPPL ~/.julia/packages/DynamicPPL/wf0dU/src/sampler.jl:83
         [30] macro expansion
            @ ~/.julia/packages/AbstractMCMC/Nw3Wn/src/sample.jl:78 [inlined]
         [31] macro expansion
            @ ~/.julia/packages/AbstractMCMC/Nw3Wn/src/logging.jl:15 [inlined]
         [32] mcmcsample(rng::Random._GLOBAL_RNG, model::DynamicPPL.Model{var"#17#18", (:X, :idx, :y, :n_gr, :predictors), (:n_gr, :predictors), (), Tuple{Matrix{Float64}, Vector{Int64}, Vector{Int64}, Int64, Int64}, Tuple{Int64, Int64}}, 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.Iterators.Pairs{Symbol, Int64, Tuple{Symbol}, NamedTuple{(:nadapts,), Tuple{Int64}}})
            @ AbstractMCMC ~/.julia/packages/AbstractMCMC/Nw3Wn/src/sample.jl:76
         [33] #sample#40
            @ ~/.julia/packages/Turing/uAz5c/src/inference/hmc.jl:140 [inlined]
         [34] macro expansion
            @ ~/.julia/packages/AbstractMCMC/Nw3Wn/src/sample.jl:280 [inlined]
         [35] (::AbstractMCMC.var"#819#threadsfor_fun#45"{UnitRange{Int64}, Bool, Base.Iterators.Pairs{Symbol, UnionAll, Tuple{Symbol}, NamedTuple{(:chain_type,), Tuple{UnionAll}}}, Int64, Vector{Any}, Vector{UInt64}, Vector{DynamicPPL.Sampler{NUTS{Turing.Core.ZygoteAD, (), AdvancedHMC.DiagEuclideanMetric}}}, Vector{DynamicPPL.Model{var"#17#18", (:X, :idx, :y, :n_gr, :predictors), (:n_gr, :predictors), (), Tuple{Matrix{Float64}, Vector{Int64}, Vector{Int64}, Int64, Int64}, Tuple{Int64, Int64}}}, Vector{Random._GLOBAL_RNG}})(onethread::Bool)
            @ AbstractMCMC ./threadingconstructs.jl:81
         [36] (::AbstractMCMC.var"#819#threadsfor_fun#45"{UnitRange{Int64}, Bool, Base.Iterators.Pairs{Symbol, UnionAll, Tuple{Symbol}, NamedTuple{(:chain_type,), Tuple{UnionAll}}}, Int64, Vector{Any}, Vector{UInt64}, Vector{DynamicPPL.Sampler{NUTS{Turing.Core.ZygoteAD, (), AdvancedHMC.DiagEuclideanMetric}}}, Vector{DynamicPPL.Model{var"#17#18", (:X, :idx, :y, :n_gr, :predictors), (:n_gr, :predictors), (), Tuple{Matrix{Float64}, Vector{Int64}, Vector{Int64}, Int64, Int64}, Tuple{Int64, Int64}}}, Vector{Random._GLOBAL_RNG}})()
            @ AbstractMCMC ./threadingconstructs.jl:48
Stacktrace:
  [1] sync_end(c::Channel{Any})
    @ Base ./task.jl:364
  [2] macro expansion
    @ ./task.jl:383 [inlined]
  [3] macro expansion
    @ ~/.julia/packages/AbstractMCMC/Nw3Wn/src/sample.jl:257 [inlined]
  [4] macro expansion
    @ ~/.julia/packages/ProgressLogging/6KXlp/src/ProgressLogging.jl:328 [inlined]
  [5] macro expansion
    @ ~/.julia/packages/AbstractMCMC/Nw3Wn/src/logging.jl:8 [inlined]
  [6] mcmcsample(rng::Random._GLOBAL_RNG, model::DynamicPPL.Model{var"#17#18", (:X, :idx, :y, :n_gr, :predictors), (:n_gr, :predictors), (), Tuple{Matrix{Float64}, Vector{Int64}, Vector{Int64}, Int64, Int64}, Tuple{Int64, Int64}}, sampler::DynamicPPL.Sampler{NUTS{Turing.Core.ZygoteAD, (), AdvancedHMC.DiagEuclideanMetric}}, ::MCMCThreads, N::Int64, nchains::Int64; progress::Bool, progressname::String, kwargs::Base.Iterators.Pairs{Symbol, UnionAll, Tuple{Symbol}, NamedTuple{(:chain_type,), Tuple{UnionAll}}})
    @ AbstractMCMC ~/.julia/packages/AbstractMCMC/Nw3Wn/src/sample.jl:251
  [7] sample(rng::Random._GLOBAL_RNG, model::DynamicPPL.Model{var"#17#18", (:X, :idx, :y, :n_gr, :predictors), (:n_gr, :predictors), (), Tuple{Matrix{Float64}, Vector{Int64}, Vector{Int64}, Int64, Int64}, Tuple{Int64, Int64}}, sampler::DynamicPPL.Sampler{NUTS{Turing.Core.ZygoteAD, (), AdvancedHMC.DiagEuclideanMetric}}, parallel::MCMCThreads, N::Int64, n_chains::Int64; chain_type::Type, progress::Bool, kwargs::Base.Iterators.Pairs{Union{}, Union{}, Tuple{}, NamedTuple{(), Tuple{}}})
    @ Turing.Inference ~/.julia/packages/Turing/uAz5c/src/inference/Inference.jl:217
  [8] sample
    @ ~/.julia/packages/Turing/uAz5c/src/inference/Inference.jl:217 [inlined]
  [9] #sample#6
    @ ~/.julia/packages/Turing/uAz5c/src/inference/Inference.jl:202 [inlined]
 [10] sample
    @ ~/.julia/packages/Turing/uAz5c/src/inference/Inference.jl:202 [inlined]
 [11] #sample#5
    @ ~/.julia/packages/Turing/uAz5c/src/inference/Inference.jl:189 [inlined]
 [12] sample(model::DynamicPPL.Model{var"#17#18", (:X, :idx, :y, :n_gr, :predictors), (:n_gr, :predictors), (), Tuple{Matrix{Float64}, Vector{Int64}, Vector{Int64}, Int64, Int64}, Tuple{Int64, Int64}}, alg::NUTS{Turing.Core.ZygoteAD, (), AdvancedHMC.DiagEuclideanMetric}, parallel::MCMCThreads, N::Int64, n_chains::Int64)
    @ Turing.Inference ~/.julia/packages/Turing/uAz5c/src/inference/Inference.jl:189
 [13] top-level scope
    @ ./timing.jl:206 [inlined]
 [14] top-level scope
    @ ~/Documents/Julia_Scripts/Turing/Turing_hierarchical.jl:0
in expression starting at /Users/storopoli/Documents/Julia_Scripts/Turing/Turing_hierarchical.jl:68