MethodError for AxisArrays.AxisArray for creating chains

[ashiklom]

I'm not sure if this is a Turing problem or a MCMCChains problem, but...when trying to fit a minimal Turing model with the following code...

using Turing
function myfun(a, b)
    x = [0:0.1:1;]
    a .* x .+ b
end
aa = 2.0
bb = 3.0
obs1 = myfun(aa, bb) .+ (rand(11) .- 0.5)
@model fitmyfun(y) = begin
    a ~ Normal(0, 10)
    b ~ Normal(0, 10)
    resid ~ InverseGamma(1, 1)
    mod = myfun(a, b)
    y ~ MvNormal(mod, sqrt(resid))
end
chain = sample(fitmyfun(obs1), NUTS(), 1000)

...I get the following error at the end of the sampling (I guess, when the chain objects are being finalized):

ERROR: LoadError: MethodError: no method matching AxisArrays.AxisArray(::Array{Float64,3}; iter=1:1:500, var=["resid", "b", "a", "lp", "n_steps", "is_accept", "acceptance_rate", "log_density", "hamiltonian_energy", "hamiltonian_energy_error", "max_hamiltonian_energy_error", "tree_depth", "numerical_error", "step_size", "nom_step_size"], chain=1:1)
Closest candidates are:
  AxisArrays.AxisArray(::AbstractArray) at /Users/ashiklom/.julia/packages/AxisArrays/yhKAp/src/core.jl:231 got unsupported keyword arguments "iter", "var", "chain"
  AxisArrays.AxisArray(::D, ::Ax) where {T, N, D<:AbstractArray{T,N}, Ax<:Tuple{Vararg{AxisArrays.Axis,N}}} at /Users/ashiklom/.julia/packages/AxisArrays/yhKAp/src/core.jl:225 got unsupported keyword arguments "iter", "var", "chain"
  AxisArrays.AxisArray(::AbstractArray, ::Tuple{Vararg{AxisArrays.Axis,N} where N}) at /Users/ashiklom/.julia/packages/AxisArrays/yhKAp/src/core.jl:221 got unsupported keyword arguments "iter", "var", "chain"
  ...
Stacktrace:
 [1] Chains(::Array{Float64,3}, ::Array{String,1}, ::NamedTuple{(:internals,),Tuple{Array{String,1}}}; start::Int64, thin::Int64, evidence::Missing, info::NamedTuple{(),Tuple{}}, sorted::Bool) at /Users/ashiklom/.julia/packages/MCMCChains/WWfCJ/src/chains.jl:84
 [2] Chains(::Array{Float64,2}, ::Array{String,1}, ::Vararg{Any,N} where N; kwargs::Base.Iterators.Pairs{Symbol,Any,Tuple{Symbol,Symbol,Symbol},NamedTuple{(:evidence, :info, :sorted),Tuple{Missing,NamedTuple{(),Tuple{}},Bool}}}) at /Users/ashiklom/.julia/packages/MCMCChains/WWfCJ/src/chains.jl:17
 [3] bundle_samples(::Random._GLOBAL_RNG, ::DynamicPPL.Model{var"##inner_function#324#4",NamedTuple{(:y,),Tuple{Array{Float64,1}}},DynamicPPL.ModelGen{(:y,),var"###fitmyfun#350",NamedTuple{(),Tuple{}}},Val{()}}, ::DynamicPPL.Sampler{NUTS{Turing.Core.ForwardDiffAD{40},(),AdvancedHMC.DiagEuclideanMetric},Turing.Inference.HMCState{DynamicPPL.VarInfo{NamedTuple{(:a, :b, :resid),Tuple{DynamicPPL.Metadata{Dict{DynamicPPL.VarName{:a},Int64},Array{Normal{Float64},1},Array{DynamicPPL.VarName{:a},1},Array{Float64,1},Array{Set{DynamicPPL.Selector},1}},DynamicPPL.Metadata{Dict{DynamicPPL.VarName{:b},Int64},Array{Normal{Float64},1},Array{DynamicPPL.VarName{:b},1},Array{Float64,1},Array{Set{DynamicPPL.Selector},1}},DynamicPPL.Metadata{Dict{DynamicPPL.VarName{:resid},Int64},Array{InverseGamma{Float64},1},Array{DynamicPPL.VarName{:resid},1},Array{Float64,1},Array{Set{DynamicPPL.Selector},1}}}},Float64},AdvancedHMC.NUTS{AdvancedHMC.MultinomialTS,AdvancedHMC.GeneralisedNoUTurn,AdvancedHMC.Leapfrog{Float64},Float64},AdvancedHMC.Adaptation.StanHMCAdaptor{AdvancedHMC.Adaptation.WelfordVar{Float64,Array{Float64,1}},AdvancedHMC.Adaptation.NesterovDualAveraging{Float64}},AdvancedHMC.PhasePoint{Array{Float64,1},AdvancedHMC.DualValue{Float64,Array{Float64,1}}}}}, ::Int64, ::Array{Turing.Inference.HamiltonianTransition{NamedTuple{(:a, :b, :resid),Tuple{Tuple{Array{Float64,1},Array{String,1}},Tuple{Array{Float64,1},Array{String,1}},Tuple{Array{Float64,1},Array{String,1}}}},NamedTuple{(:n_steps, :is_accept, :acceptance_rate, :log_density, :hamiltonian_energy, :hamiltonian_energy_error, :max_hamiltonian_energy_error, :tree_depth, :numerical_error, :step_size, :nom_step_size),Tuple{Int64,Bool,Float64,Float64,Float64,Float64,Float64,Int64,Bool,Float64,Float64}},Float64},1}, ::Type{Chains}; discard_adapt::Bool, save_state::Bool, kwargs::Base.Iterators.Pairs{Union{},Union{},Tuple{},NamedTuple{(),Tuple{}}}) at /Users/ashiklom/.julia/packages/Turing/q3coT/src/inference/Inference.jl:365
 [4] bundle_samples(::Random._GLOBAL_RNG, ::DynamicPPL.Model{var"##inner_function#324#4",NamedTuple{(:y,),Tuple{Array{Float64,1}}},DynamicPPL.ModelGen{(:y,),var"###fitmyfun#350",NamedTuple{(),Tuple{}}},Val{()}}, ::DynamicPPL.Sampler{NUTS{Turing.Core.ForwardDiffAD{40},(),AdvancedHMC.DiagEuclideanMetric},Turing.Inference.HMCState{DynamicPPL.VarInfo{NamedTuple{(:a, :b, :resid),Tuple{DynamicPPL.Metadata{Dict{DynamicPPL.VarName{:a},Int64},Array{Normal{Float64},1},Array{DynamicPPL.VarName{:a},1},Array{Float64,1},Array{Set{DynamicPPL.Selector},1}},DynamicPPL.Metadata{Dict{DynamicPPL.VarName{:b},Int64},Array{Normal{Float64},1},Array{DynamicPPL.VarName{:b},1},Array{Float64,1},Array{Set{DynamicPPL.Selector},1}},DynamicPPL.Metadata{Dict{DynamicPPL.VarName{:resid},Int64},Array{InverseGamma{Float64},1},Array{DynamicPPL.VarName{:resid},1},Array{Float64,1},Array{Set{DynamicPPL.Selector},1}}}},Float64},AdvancedHMC.NUTS{AdvancedHMC.MultinomialTS,AdvancedHMC.GeneralisedNoUTurn,AdvancedHMC.Leapfrog{Float64},Float64},AdvancedHMC.Adaptation.StanHMCAdaptor{AdvancedHMC.Adaptation.WelfordVar{Float64,Array{Float64,1}},AdvancedHMC.Adaptation.NesterovDualAveraging{Float64}},AdvancedHMC.PhasePoint{Array{Float64,1},AdvancedHMC.DualValue{Float64,Array{Float64,1}}}}}, ::Int64, ::Array{Turing.Inference.HamiltonianTransition{NamedTuple{(:a, :b, :resid),Tuple{Tuple{Array{Float64,1},Array{String,1}},Tuple{Array{Float64,1},Array{String,1}},Tuple{Array{Float64,1},Array{String,1}}}},NamedTuple{(:n_steps, :is_accept, :acceptance_rate, :log_density, :hamiltonian_energy, :hamiltonian_energy_error, :max_hamiltonian_energy_error, :tree_depth, :numerical_error, :step_size, :nom_step_size),Tuple{Int64,Bool,Float64,Float64,Float64,Float64,Float64,Int64,Bool,Float64,Float64}},Float64},1}, ::Type{Chains}) at /Users/ashiklom/.julia/packages/Turing/q3coT/src/inference/Inference.jl:326
 [5] mcmcsample(::Random._GLOBAL_RNG, ::DynamicPPL.Model{var"##inner_function#324#4",NamedTuple{(:y,),Tuple{Array{Float64,1}}},DynamicPPL.ModelGen{(:y,),var"###fitmyfun#350",NamedTuple{(),Tuple{}}},Val{()}}, ::DynamicPPL.Sampler{NUTS{Turing.Core.ForwardDiffAD{40},(),AdvancedHMC.DiagEuclideanMetric},Turing.Inference.HMCState{DynamicPPL.VarInfo{NamedTuple{(:a, :b, :resid),Tuple{DynamicPPL.Metadata{Dict{DynamicPPL.VarName{:a},Int64},Array{Normal{Float64},1},Array{DynamicPPL.VarName{:a},1},Array{Float64,1},Array{Set{DynamicPPL.Selector},1}},DynamicPPL.Metadata{Dict{DynamicPPL.VarName{:b},Int64},Array{Normal{Float64},1},Array{DynamicPPL.VarName{:b},1},Array{Float64,1},Array{Set{DynamicPPL.Selector},1}},DynamicPPL.Metadata{Dict{DynamicPPL.VarName{:resid},Int64},Array{InverseGamma{Float64},1},Array{DynamicPPL.VarName{:resid},1},Array{Float64,1},Array{Set{DynamicPPL.Selector},1}}}},Float64},AdvancedHMC.NUTS{AdvancedHMC.MultinomialTS,AdvancedHMC.GeneralisedNoUTurn,AdvancedHMC.Leapfrog{Float64},Float64},AdvancedHMC.Adaptation.StanHMCAdaptor{AdvancedHMC.Adaptation.WelfordVar{Float64,Array{Float64,1}},AdvancedHMC.Adaptation.NesterovDualAveraging{Float64}},AdvancedHMC.PhasePoint{Array{Float64,1},AdvancedHMC.DualValue{Float64,Array{Float64,1}}}}}, ::Int64; progress::Bool, progressname::String, callback::AbstractMCMC.var"#21#26"{AbstractMCMC.var"#21#22#27"}, chain_type::Type{T} where T, kwargs::Base.Iterators.Pairs{Union{},Union{},Tuple{},NamedTuple{(),Tuple{}}}) at /Users/ashiklom/.julia/packages/AbstractMCMC/MPGEV/src/sample.jl:106
 [6] sample(::Random._GLOBAL_RNG, ::DynamicPPL.Model{var"##inner_function#324#4",NamedTuple{(:y,),Tuple{Array{Float64,1}}},DynamicPPL.ModelGen{(:y,),var"###fitmyfun#350",NamedTuple{(),Tuple{}}},Val{()}}, ::DynamicPPL.Sampler{NUTS{Turing.Core.ForwardDiffAD{40},(),AdvancedHMC.DiagEuclideanMetric},Turing.Inference.HMCState{DynamicPPL.VarInfo{NamedTuple{(:a, :b, :resid),Tuple{DynamicPPL.Metadata{Dict{DynamicPPL.VarName{:a},Int64},Array{Normal{Float64},1},Array{DynamicPPL.VarName{:a},1},Array{Float64,1},Array{Set{DynamicPPL.Selector},1}},DynamicPPL.Metadata{Dict{DynamicPPL.VarName{:b},Int64},Array{Normal{Float64},1},Array{DynamicPPL.VarName{:b},1},Array{Float64,1},Array{Set{DynamicPPL.Selector},1}},DynamicPPL.Metadata{Dict{DynamicPPL.VarName{:resid},Int64},Array{InverseGamma{Float64},1},Array{DynamicPPL.VarName{:resid},1},Array{Float64,1},Array{Set{DynamicPPL.Selector},1}}}},Float64},AdvancedHMC.NUTS{AdvancedHMC.MultinomialTS,AdvancedHMC.GeneralisedNoUTurn,AdvancedHMC.Leapfrog{Float64},Float64},AdvancedHMC.Adaptation.StanHMCAdaptor{AdvancedHMC.Adaptation.WelfordVar{Float64,Array{Float64,1}},AdvancedHMC.Adaptation.NesterovDualAveraging{Float64}},AdvancedHMC.PhasePoint{Array{Float64,1},AdvancedHMC.DualValue{Float64,Array{Float64,1}}}}}, ::Int64; kwargs::Base.Iterators.Pairs{Symbol,Any,Tuple{Symbol,Symbol},NamedTuple{(:chain_type, :progress),Tuple{UnionAll,Bool}}}) at /Users/ashiklom/.julia/packages/AbstractMCMC/MPGEV/src/sample.jl:19
 [7] sample(::Random._GLOBAL_RNG, ::DynamicPPL.Model{var"##inner_function#324#4",NamedTuple{(:y,),Tuple{Array{Float64,1}}},DynamicPPL.ModelGen{(:y,),var"###fitmyfun#350",NamedTuple{(),Tuple{}}},Val{()}}, ::NUTS{Turing.Core.ForwardDiffAD{40},(),AdvancedHMC.DiagEuclideanMetric}, ::Int64; chain_type::Type{T} where T, resume_from::Nothing, progress::Bool, kwargs::Base.Iterators.Pairs{Union{},Union{},Tuple{},NamedTuple{(),Tuple{}}}) at /Users/ashiklom/.julia/packages/Turing/q3coT/src/inference/Inference.jl:155
 [8] sample at /Users/ashiklom/.julia/packages/Turing/q3coT/src/inference/Inference.jl:154 [inlined]
 [9] #sample#1 at /Users/ashiklom/.julia/packages/Turing/q3coT/src/inference/Inference.jl:141 [inlined]
 [10] sample(::DynamicPPL.Model{var"##inner_function#324#4",NamedTuple{(:y,),Tuple{Array{Float64,1}}},DynamicPPL.ModelGen{(:y,),var"###fitmyfun#350",NamedTuple{(),Tuple{}}},Val{()}}, ::NUTS{Turing.Core.ForwardDiffAD{40},(),AdvancedHMC.DiagEuclideanMetric}, ::Int64) at /Users/ashiklom/.julia/packages/Turing/q3coT/src/inference/Inference.jl:141
 [11] top-level scope at /Users/ashiklom/projects/other/julia-experimentation/turing-mre/turing-mre.jl:20
 [12] include(::String) at ./client.jl:439
 [13] top-level scope at REPL[3]:1
in expression starting at /Users/ashiklom/projects/other/julia-experimentation/turing-mre/turing-mre.jl:20

I get the same error using the latest stable release (v0.10.1) and the GitHub master branch.

[ashiklom]

On closer investigation, this is definitely a MCMCChains problem. Opening an issue there and closing this one.

[devmotion]

I assume you do not use AxisArrays 0.4.2 (I'm not sure why this is happening, also the latest CI tests in Turing used AxisArrays 0.4.0 and hence failed - but this error only occurred for the latest commit on Turing master, prior to that everything worked as expected). Can you check if that's the case for you?

You can check which package is holding AxisArrays back by running

julia> ] add AxisArrays@0.4.2

[ashiklom]

Fixed by upgrading to AxisArrays@0.4.2 (or, indirectly, by TuringLang/MCMCChains.jl#213).