StaticArrays throw UndefRefError

[penelopeysm]

using DynamicPPL, Distributions, StaticArrays
@model function f3()
    x = @SVector zeros(Float64, 3)
    x .~ Beta(2, 2)
end
m3 = f3()
VarInfo(m3)

ERROR: UndefRefError: access to undefined reference
Stacktrace:
  [1] getindex
    @ ./essentials.jl:917 [inlined]
  [2] getindex
    @ ~/.julia/packages/StaticArrays/IZgIP/src/SizedArray.jl:92 [inlined]
  [3] macro expansion
    @ ~/.julia/packages/StaticArrays/IZgIP/src/deque.jl:193 [inlined]
  [4] _setindex
    @ ~/.julia/packages/StaticArrays/IZgIP/src/deque.jl:186 [inlined]
  [5] setindex
    @ ~/.julia/packages/StaticArrays/IZgIP/src/deque.jl:185 [inlined]
  [6] _setindex
    @ ~/.julia/packages/BangBang/9Z60G/ext/BangBangStaticArraysExt.jl:27 [inlined]
  [7] may
    @ ~/.julia/packages/BangBang/9Z60G/src/core.jl:11 [inlined]
  [8] setindex!!
    @ ~/.julia/packages/BangBang/9Z60G/src/base.jl:478 [inlined]
  [9] #setindex!!#23
    @ ~/ppl/dppl/src/varnamedtuple/partial_array.jl:584 [inlined]
 [10] setindex!!
    @ ~/ppl/dppl/src/varnamedtuple/partial_array.jl:541 [inlined]
 [11] make_leaf_singleindex(value::VectorValue{…}, coptic::AbstractPPL.Index{…}, template::SVector{…})
    @ DynamicPPL.VarNamedTuples ~/ppl/dppl/src/varnamedtuple/getset.jl:342
 [12] make_leaf
    @ ~/ppl/dppl/src/varnamedtuple/getset.jl:291 [inlined]
 [13] #_setindex_optic!!#35
    @ ~/ppl/dppl/src/varnamedtuple/getset.jl:224 [inlined]
 [14] _setindex_optic!!
    @ ~/ppl/dppl/src/varnamedtuple/getset.jl:205 [inlined]
 [15] templated_setindex!!
    @ ~/ppl/dppl/src/varnamedtuple/vnt.jl:115 [inlined]

[penelopeysm]

As for an explanation of how I came across this: I was investigating the performance of models like

@model function f()
    x = zeros(100)
    for i in eachindex(x)
        x[i] ~ Beta(2, 2)
    end
end

The problem with this model is that it is actually impossible to statically infer whether the resulting VarInfo will have any keys of the form x[i]. The reason is because it is not known until compile time whether eachindex(x) is empty (in which case the VarInfo will have no keys), or whether it is nonempty (in which case the VarInfo will have some keys).

julia> @code_warntype VarInfo(f())
MethodInstance for VarInfo(::Model{typeof(f), (), (), (), Tuple{}, Tuple{}, DefaultContext, false})
  from VarInfo(model::Model) @ DynamicPPL ~/ppl/dppl/src/varinfo.jl:150
Arguments
  #self#::Type{VarInfo}
  model::Core.Const(Model{typeof(f), (), (), (), Tuple{}, Tuple{}, DefaultContext, false}(Main.f, NamedTuple(), NamedTuple(), DefaultContext()))
Body::Union{VarInfo{false, VarNamedTuple{(), Tuple{}}, DynamicPPL.AccumulatorTuple{3, @NamedTuple{LogPrior::LogPriorAccumulator{Float64}, LogJacobian::LogJacobianAccumulator{Float64}, LogLikelihood::LogLikelihoodAccumulator{Float64}}}}, VarInfo{false, VarNamedTuple{(:x,), Tuple{DynamicPPL.VarNamedTuples.PartialArray{VectorValue{Vector{Float64}, DynamicPPL.UnwrapSingletonTransform{Tuple{Int64}}, Tuple{}}, 1, Vector{VectorValue{Vector{Float64}, DynamicPPL.UnwrapSingletonTransform{Tuple{Int64}}, Tuple{}}}, Vector{Bool}}}}, DynamicPPL.AccumulatorTuple{3, @NamedTuple{LogPrior::LogPriorAccumulator{Float64}, LogJacobian::LogJacobianAccumulator{Float64}, LogLikelihood::LogLikelihoodAccumulator{Float64}}}}}
1 ─ %1 = DynamicPPL.InitFromPrior()::Core.Const(InitFromPrior())
│   %2 = (#self#)(model, %1)::Union{VarInfo{false, VarNamedTuple{(), Tuple{}}, DynamicPPL.AccumulatorTuple{3, @NamedTuple{LogPrior::LogPriorAccumulator{Float64}, LogJacobian::LogJacobianAccumulator{Float64}, LogLikelihood::LogLikelihoodAccumulator{Float64}}}}, VarInfo{false, VarNamedTuple{(:x,), Tuple{DynamicPPL.VarNamedTuples.PartialArray{VectorValue{Vector{Float64}, DynamicPPL.UnwrapSingletonTransform{Tuple{Int64}}, Tuple{}}, 1, Vector{VectorValue{Vector{Float64}, DynamicPPL.UnwrapSingletonTransform{Tuple{Int64}}, Tuple{}}}, Vector{Bool}}}}, DynamicPPL.AccumulatorTuple{3, @NamedTuple{LogPrior::LogPriorAccumulator{Float64}, LogJacobian::LogJacobianAccumulator{Float64}, LogLikelihood::LogLikelihoodAccumulator{Float64}}}}}
└──      return %2

A more readable form of that return type is:

Union{

  # Empty VarInfo.
  VarInfo{
    false,
    VarNamedTuple{(), Tuple{}},
    DynamicPPL.AccumulatorTuple{3, @NamedTuple{LogPrior::LogPriorAccumulator{Float64}, LogJacobian::LogJacobianAccumulator{Float64}, LogLikelihood::LogLikelihoodAccumulator{Float64}}}
  },

  # VarInfo with some keys `x[i]`.
  VarInfo{
    false,
    VarNamedTuple{(:x,), Tuple{DynamicPPL.VarNamedTuples.PartialArray{VectorValue{Vector{Float64}, DynamicPPL.UnwrapSingletonTransform{Tuple{Int64}}, Tuple{}}, 1, Vector{VectorValue{Vector{Float64}, DynamicPPL.UnwrapSingletonTransform{Tuple{Int64}}, Tuple{}}}, Vector{Bool}}}},
    DynamicPPL.AccumulatorTuple{3, @NamedTuple{LogPrior::LogPriorAccumulator{Float64}, LogJacobian::LogJacobianAccumulator{Float64}, LogLikelihood::LogLikelihoodAccumulator{Float64}}}
  }

}

Note that this is a general 'problem' with iteration in Julia so it's not something that we can solve. For example:

function f(x)
    y = 1.0
    for i in eachindex(x)
        y = :a
    end
    return y
end

@code_warntype f(zeros(3)) # --> Union{Float64,Symbol}

using StaticArrays
@code_warntype f(SVector(0.0, 0.0, 0.0)) # --> Symbol

using FixedSizeArrays
@code_warntype f(FixedSizeArray(zeros(3))) # --> Union{Float64,Symbol}

I thought that since StaticArrays work here, it might also allow the original model to be inferred, but it errored.

FixedSizeArrays work in the model, but still fails to infer.

[penelopeysm]

Minimised to JuliaArrays/StaticArrays.jl#1330. It could be a BangBang thing too as it seems questionable to use setindex when setindex! would have worked fine.

[penelopeysm]

Opened a BangBang PR that would fix this JuliaFolds2/BangBang.jl#45, with it the model above runs fine

VarInfo
 ├─ values
 │  VarNamedTuple
 │  └─ x => PartialArray size=(3,) data::SizedVector{3, VectorValue{Vector{Float64}, DynamicPPL.UnwrapSingletonTransform{Tuple{Int64}}, Tuple{}}, Vector{VectorValue{Vector{Float64}, DynamicPPL.UnwrapSingletonTransform{Tuple{Int64}}, Tuple{}}}}
 │          ├─ (1,) => VectorValue{Vector{Float64}, DynamicPPL.UnwrapSingletonTransform{Tuple{Int64}}, Tuple{}}([0.36038245345012115], DynamicPPL.UnwrapSingletonTransform{Tuple{Int64}}((1,)), ())
 │          ├─ (2,) => VectorValue{Vector{Float64}, DynamicPPL.UnwrapSingletonTransform{Tuple{Int64}}, Tuple{}}([0.28555079817134305], DynamicPPL.UnwrapSingletonTransform{Tuple{Int64}}((1,)), ())
 │          └─ (3,) => VectorValue{Vector{Float64}, DynamicPPL.UnwrapSingletonTransform{Tuple{Int64}}, Tuple{}}([0.879706978496964], DynamicPPL.UnwrapSingletonTransform{Tuple{Int64}}((1,)), ())
 └─ accs
    AccumulatorTuple with 3 accumulators
    ├─ LogPrior => LogPriorAccumulator(0.0722343105001233)
    ├─ LogJacobian => LogJacobianAccumulator(0.0)
    └─ LogLikelihood => LogLikelihoodAccumulator(0.0)