Turing models, vectors, and type stability?

Hi,
I’m puzzled by a type instability issue in the Turing model below. When the model defines the vector p2 before iterating, the model is type stable. When using the [...] notation, the model is unstable.

using LinearAlgebra
using Random
using Turing

@model function tmodel_v1(x, mat)
    Nx, _ = size(mat)
    p ~ filldist(Normal(0.0, 1.0), Nx)
    σ² ~ truncated(Normal(0.0, 0.01); lower=0)

    p2 = [dot(p, m) for m ∈ eachslice(mat, dims=2)]

    x ~ MvNormal(p2, σ² * I)
end

@model function tmodel_v2(x, mat)
    Nx, Ny = size(mat)

    p ~ filldist(Normal(0.0, 1.0), Nx)
    σ² ~ truncated(Normal(0.0, 0.01); lower=0)

    p2 = zeros(Ny)
    for ix ∈ axes(mat, 2)
        p2[ix] = dot(p, mat[:, ix])
    end

    x ~ MvNormal(p2, σ² * I)
end

x = randn(4)
A = randn(3, 4)

model_unstable = tmodel_v1(x, A)
model_stable = tmodel_v2(x, A)

@code_warntype model_unstable.f(
    model_unstable,
    Turing.VarInfo(model_unstable),
    Turing.SamplingContext(
        Random.GLOBAL_RNG, Turing.SampleFromPrior(), Turing.DefaultContext(),
    ),
    model_unstable.args...,
)

@code_warntype model_stable.f(
    model_unstable,
    Turing.VarInfo(model_unstable),
    Turing.SamplingContext(
        Random.GLOBAL_RNG, Turing.SampleFromPrior(), Turing.DefaultContext(),
    ),
    model_unstable.args...,
)

  from tmodel_v1(__model__::DynamicPPL.Model, __varinfo__::DynamicPPL.AbstractVarInfo, __context__::AbstractPPL.AbstractContext, x, mat) @ Main ~/julia_envs/bayes_fun/src/test_typestability.jl:7
Arguments
  #self#::Core.Const(tmodel_v1)
  __model__::DynamicPPL.Model{typeof(tmodel_v1), (:x, :mat), (), (), Tuple{Vector{Float64}, Matrix{Float64}}, Tuple{}, DynamicPPL.DefaultContext}
  __varinfo__@_3::DynamicPPL.TypedVarInfo{@NamedTuple{p::DynamicPPL.Metadata{Dict{AbstractPPL.VarName{:p, typeof(identity)}, Int64}, Vector{DistributionsAD.TuringScalMvNormal{Vector{Float64}, Float64}}, Vector{AbstractPPL.VarName{:p, typeof(identity)}}, Vector{Float64}, Vector{Set{DynamicPPL.Selector}}}, σ²::DynamicPPL.Metadata{Dict{AbstractPPL.VarName{:σ², typeof(identity)}, Int64}, Vector{Truncated{Normal{Float64}, Continuous, Float64, Float64, Nothing}}, Vector{AbstractPPL.VarName{:σ², typeof(identity)}}, Vector{Float64}, Vector{Set{DynamicPPL.Selector}}}}, Float64}
  __context__::Core.Const(DynamicPPL.SamplingContext{DynamicPPL.SampleFromPrior, DynamicPPL.DefaultContext, Random._GLOBAL_RNG}(Random._GLOBAL_RNG(), DynamicPPL.SampleFromPrior(), DynamicPPL.DefaultContext()))
  x@_5::Vector{Float64}
  mat::Matrix{Float64}
Locals
  @_7::Int64
  @_8::Int64
  #11::var"#11#12"
  @_10::Union{}
  @_11::Int64
  @_12::Union{}
  @_13::Int64
  @_14::Int64
  retval#240::Any
  value#237::Vector{Float64}
  value#239::Any
  isassumption#236::Bool
  vn#235::AbstractPPL.VarName
  dist#238::Any
  p2::Vector
  value#232::Union{}
  σ²::Float64
  value#234::Float64
  isassumption#231::Bool
  vn#230::AbstractPPL.VarName{:σ², typeof(identity)}
  dist#233::Truncated{Normal{Float64}, Continuous, Float64, Float64, Nothing}
  value#227::Union{}
  p@_29::Core.Box
  value#229::Vector{Float64}
  isassumption#226::Bool
  vn#225::AbstractPPL.VarName{:p, typeof(identity)}
  dist#228::DistributionsAD.TuringScalMvNormal{Vector{Float64}, Float64}
  Nx::Int64
  p@_35::Union{}
  p@_36::Union{}
  __varinfo__@_37::DynamicPPL.TypedVarInfo{@NamedTuple{p::DynamicPPL.Metadata{Dict{AbstractPPL.VarName{:p, typeof(identity)}, Int64}, Vector{DistributionsAD.TuringScalMvNormal{Vector{Float64}, Float64}}, Vector{AbstractPPL.VarName{:p, typeof(identity)}}, Vector{Float64}, Vector{Set{DynamicPPL.Selector}}}, σ²::DynamicPPL.Metadata{Dict{AbstractPPL.VarName{:σ², typeof(identity)}, Int64}, Vector{Truncated{Normal{Float64}, Continuous, Float64, Float64, Nothing}}, Vector{AbstractPPL.VarName{:σ², typeof(identity)}}, Vector{Float64}, Vector{Set{DynamicPPL.Selector}}}}, Float64}
  x@_38::Any
  @_39::Bool
  @_40::Bool
  @_41::Bool
  @_42::Bool
  @_43::Bool
  @_44::Bool
  @_45::Any
  @_46::Any
Body::Tuple{Any, DynamicPPL.TypedVarInfo{@NamedTuple{p::DynamicPPL.Metadata{Dict{AbstractPPL.VarName{:p, typeof(identity)}, Int64}, Vector{DistributionsAD.TuringScalMvNormal{Vector{Float64}, Float64}}, Vector{AbstractPPL.VarName{:p, typeof(identity)}}, Vector{Float64}, Vector{Set{DynamicPPL.Selector}}}, σ²::DynamicPPL.Metadata{Dict{AbstractPPL.VarName{:σ², typeof(identity)}, Int64}, Vector{Truncated{Normal{Float64}, Continuous, Float64, Float64, Nothing}}, Vector{AbstractPPL.VarName{:σ², typeof(identity)}}, Vector{Float64}, Vector{Set{DynamicPPL.Selector}}}}, Float64}}
1 ──        (x@_38 = x@_5)
│           (__varinfo__@_37 = __varinfo__@_3)
│           Core.NewvarNode(:(@_7))
│           Core.NewvarNode(:(@_8))
│           Core.NewvarNode(:(#11))
│           Core.NewvarNode(:(@_10))
│           Core.NewvarNode(:(@_11))
│           Core.NewvarNode(:(@_12))
│           Core.NewvarNode(:(@_13))
│           Core.NewvarNode(:(retval#240))
│           Core.NewvarNode(:(value#237))
│           Core.NewvarNode(:(value#239))
│           Core.NewvarNode(:(isassumption#236))
│           Core.NewvarNode(:(vn#235))
│           Core.NewvarNode(:(dist#238))
│           Core.NewvarNode(:(p2))
│           Core.NewvarNode(:(value#232))
│           Core.NewvarNode(:(σ²))
│           Core.NewvarNode(:(value#234))
│           Core.NewvarNode(:(isassumption#231))
│           Core.NewvarNode(:(vn#230))
│           Core.NewvarNode(:(dist#233))
│           Core.NewvarNode(:(value#227))
│           (p@_29 = Core.Box())
│           Core.NewvarNode(:(value#229))
│           Core.NewvarNode(:(isassumption#226))
│    %27  = Main.size(mat)::Tuple{Int64, Int64}
│    %28  = Base.indexed_iterate(%27, 1)::Core.PartialStruct(Tuple{Int64, Int64}, Any[Int64, Core.Const(2)])
│           (Nx = Core.getfield(%28, 1))
│           (@_14 = Core.getfield(%28, 2))
│    %31  = Base.indexed_iterate(%27, 2, @_14::Core.Const(2))::Core.PartialStruct(Tuple{Int64, Int64}, Any[Int64, Core.Const(3)])
│           Core.getfield(%31, 1)
│    %33  = Main.Normal(0.0, 1.0)::Core.Const(Normal{Float64}(μ=0.0, σ=1.0))
│           (dist#228 = Main.filldist(%33, Nx))
│    %35  = Core.apply_type(AbstractPPL.VarName, :p)::Core.Const(AbstractPPL.VarName{:p})
│    %36  = (%35)()::Core.Const(p)
│           (vn#225 = (DynamicPPL.resolve_varnames)(%36, dist#228::Core.PartialStruct(DistributionsAD.TuringScalMvNormal{Vector{Float64}, Float64}, Any[Vector{Float64}, Core.Const(1.0)])))
│    %38  = (DynamicPPL.contextual_isassumption)(__context__, vn#225)::Core.Const(true)
└───        goto #8 if not %38
2 ── %40  = (DynamicPPL.inargnames)(vn#225, __model__)::Core.Const(false)
│    %41  = !%40::Core.Const(true)
└───        goto #4 if not %41
3 ──        goto #5
4 ──        Core.Const(:((DynamicPPL.inmissings)(vn#225, __model__)))
└───        Core.Const(:(goto %48 if not %44))
5 ┄─        (@_40 = true)
└───        goto #7
6 ──        Core.Const(:(Core.isdefined(p@_29, :contents)))
│           Core.Const(:(goto %51 if not %48))
│           Core.Const(:(goto %53))
│           Core.Const(Core.NewvarNode(:(p@_35)))
│           Core.Const(:(p@_35))
│           Core.Const(:(Core.getfield(p@_29, :contents)))
└───        Core.Const(:(@_40 = %53 === Main.missing))
7 ┄─        (@_39 = @_40::Core.Const(true))
└───        goto #9
8 ──        Core.Const(:(@_39 = false))
9 ┄─        (isassumption#226 = @_39::Core.Const(true))
│    %59  = (DynamicPPL.contextual_isfixed)(__context__, vn#225)::Core.Const(false)
└───        goto #11 if not %59
10 ─        Core.Const(:((DynamicPPL.getfixed_nested)(__context__, vn#225)))
│           Core.Const(:(Core.setfield!(p@_29, :contents, %61)))
└───        Core.Const(:(goto %99))
11 ┄        goto #13 if not isassumption#226::Core.Const(true)
12 ─ %65  = Core.tuple(__context__)::Core.Const((DynamicPPL.SamplingContext{DynamicPPL.SampleFromPrior, DynamicPPL.DefaultContext, Random._GLOBAL_RNG}(Random._GLOBAL_RNG(), DynamicPPL.SampleFromPrior(), DynamicPPL.DefaultContext()),))
│    %66  = (DynamicPPL.check_tilde_rhs)(dist#228::Core.PartialStruct(DistributionsAD.TuringScalMvNormal{Vector{Float64}, Float64}, Any[Vector{Float64}, Core.Const(1.0)]))::Core.PartialStruct(DistributionsAD.TuringScalMvNormal{Vector{Float64}, Float64}, Any[Vector{Float64}, Core.Const(1.0)])
│    %67  = (DynamicPPL.unwrap_right_vn)(%66, vn#225)::Core.PartialStruct(Tuple{DistributionsAD.TuringScalMvNormal{Vector{Float64}, Float64}, AbstractPPL.VarName{:p, typeof(identity)}}, Any[Core.PartialStruct(DistributionsAD.TuringScalMvNormal{Vector{Float64}, Float64}, Any[Vector{Float64}, Core.Const(1.0)]), AbstractPPL.VarName{:p, typeof(identity)}])
│    %68  = Core.tuple(__varinfo__@_37)::Tuple{DynamicPPL.TypedVarInfo{@NamedTuple{p::DynamicPPL.Metadata{Dict{AbstractPPL.VarName{:p, typeof(identity)}, Int64}, Vector{DistributionsAD.TuringScalMvNormal{Vector{Float64}, Float64}}, Vector{AbstractPPL.VarName{:p, typeof(identity)}}, Vector{Float64}, Vector{Set{DynamicPPL.Selector}}}, σ²::DynamicPPL.Metadata{Dict{AbstractPPL.VarName{:σ², typeof(identity)}, Int64}, Vector{Truncated{Normal{Float64}, Continuous, Float64, Float64, Nothing}}, Vector{AbstractPPL.VarName{:σ², typeof(identity)}}, Vector{Float64}, Vector{Set{DynamicPPL.Selector}}}}, Float64}}
│    %69  = Core._apply_iterate(Base.iterate, DynamicPPL.tilde_assume!!, %65, %67, %68)::Tuple{Vector{Float64}, DynamicPPL.TypedVarInfo{@NamedTuple{p::DynamicPPL.Metadata{Dict{AbstractPPL.VarName{:p, typeof(identity)}, Int64}, Vector{DistributionsAD.TuringScalMvNormal{Vector{Float64}, Float64}}, Vector{AbstractPPL.VarName{:p, typeof(identity)}}, Vector{Float64}, Vector{Set{DynamicPPL.Selector}}}, σ²::DynamicPPL.Metadata{Dict{AbstractPPL.VarName{:σ², typeof(identity)}, Int64}, Vector{Truncated{Normal{Float64}, Continuous, Float64, Float64, Nothing}}, Vector{AbstractPPL.VarName{:σ², typeof(identity)}}, Vector{Float64}, Vector{Set{DynamicPPL.Selector}}}}, Float64}}
│    %70  = Base.indexed_iterate(%69, 1)::Core.PartialStruct(Tuple{Vector{Float64}, Int64}, Any[Vector{Float64}, Core.Const(2)])
│           (value#229 = Core.getfield(%70, 1))
│           (@_13 = Core.getfield(%70, 2))
│    %73  = Base.indexed_iterate(%69, 2, @_13::Core.Const(2))::Core.PartialStruct(Tuple{DynamicPPL.TypedVarInfo{@NamedTuple{p::DynamicPPL.Metadata{Dict{AbstractPPL.VarName{:p, typeof(identity)}, Int64}, Vector{DistributionsAD.TuringScalMvNormal{Vector{Float64}, Float64}}, Vector{AbstractPPL.VarName{:p, typeof(identity)}}, Vector{Float64}, Vector{Set{DynamicPPL.Selector}}}, σ²::DynamicPPL.Metadata{Dict{AbstractPPL.VarName{:σ², typeof(identity)}, Int64}, Vector{Truncated{Normal{Float64}, Continuous, Float64, Float64, Nothing}}, Vector{AbstractPPL.VarName{:σ², typeof(identity)}}, Vector{Float64}, Vector{Set{DynamicPPL.Selector}}}}, Float64}, Int64}, Any[DynamicPPL.TypedVarInfo{@NamedTuple{p::DynamicPPL.Metadata{Dict{AbstractPPL.VarName{:p, typeof(identity)}, Int64}, Vector{DistributionsAD.TuringScalMvNormal{Vector{Float64}, Float64}}, Vector{AbstractPPL.VarName{:p, typeof(identity)}}, Vector{Float64}, Vector{Set{DynamicPPL.Selector}}}, σ²::DynamicPPL.Metadata{Dict{AbstractPPL.VarName{:σ², typeof(identity)}, Int64}, Vector{Truncated{Normal{Float64}, Continuous, Float64, Float64, Nothing}}, Vector{AbstractPPL.VarName{:σ², typeof(identity)}}, Vector{Float64}, Vector{Set{DynamicPPL.Selector}}}}, Float64}, Core.Const(3)])
│           (__varinfo__@_37 = Core.getfield(%73, 1))
│    %75  = value#229::Vector{Float64}
│           Core.setfield!(p@_29, :contents, %75)
│           value#229
└───        goto #14
13 ─        Core.Const(:((DynamicPPL.inargnames)(vn#225, __model__)))
│           Core.Const(:(!%79))
│           Core.Const(:(goto %84 if not %80))
│           Core.Const(:((DynamicPPL.getconditioned_nested)(__context__, vn#225)))
│           Core.Const(:(Core.setfield!(p@_29, :contents, %82)))
│           Core.Const(:((DynamicPPL.check_tilde_rhs)(dist#228)))
│           Core.Const(:(Core.isdefined(p@_29, :contents)))
│           Core.Const(:(goto %88 if not %85))
│           Core.Const(:(goto %90))
│           Core.Const(Core.NewvarNode(:(p@_36)))
│           Core.Const(:(p@_36))
│           Core.Const(:(Core.getfield(p@_29, :contents)))
│           Core.Const(:(vn#225))
│           Core.Const(:((DynamicPPL.tilde_observe!!)(__context__, %84, %90, %91, __varinfo__@_37)))
│           Core.Const(:(Base.indexed_iterate(%92, 1)))
│           Core.Const(:(value#227 = Core.getfield(%93, 1)))
│           Core.Const(:(@_12 = Core.getfield(%93, 2)))
│           Core.Const(:(Base.indexed_iterate(%92, 2, @_12)))
│           Core.Const(:(__varinfo__@_37 = Core.getfield(%96, 1)))
└───        Core.Const(:(value#227))
14 ┄ %99  = Main.Normal(0.0, 0.01)::Core.Const(Normal{Float64}(μ=0.0, σ=0.01))
│    %100 = (:lower,)::Core.Const((:lower,))
│    %101 = Core.apply_type(Core.NamedTuple, %100)::Core.Const(NamedTuple{(:lower,)})
│    %102 = Core.tuple(0)::Core.Const((0,))
│    %103 = (%101)(%102)::Core.Const((lower = 0,))
│           (dist#233 = Core.kwcall(%103, Main.truncated, %99))
│    %105 = Core.apply_type(AbstractPPL.VarName, :σ²)::Core.Const(AbstractPPL.VarName{:σ²})
│    %106 = (%105)()::Core.Const(σ²)
│           (vn#230 = (DynamicPPL.resolve_varnames)(%106, dist#233::Core.PartialStruct(Truncated{Normal{Float64}, Continuous, Float64, Float64, Nothing}, Any[Core.Const(Normal{Float64}(μ=0.0, σ=0.01)), Core.Const(0.0), Core.Const(nothing), Float64, Float64, Core.Const(1.0), Float64, Float64])))
│    %108 = (DynamicPPL.contextual_isassumption)(__context__, vn#230)::Core.Const(true)
└───        goto #21 if not %108
15 ─ %110 = (DynamicPPL.inargnames)(vn#230, __model__)::Core.Const(false)
│    %111 = !%110::Core.Const(true)
└───        goto #17 if not %111
16 ─        goto #18
17 ─        Core.Const(:((DynamicPPL.inmissings)(vn#230, __model__)))
└───        Core.Const(:(goto %118 if not %114))
18 ┄        (@_42 = true)
└───        goto #20
19 ─        Core.Const(:(@_42 = σ² === Main.missing))
20 ┄        (@_41 = @_42::Core.Const(true))
└───        goto #22
21 ─        Core.Const(:(@_41 = false))
22 ┄        (isassumption#231 = @_41::Core.Const(true))
│    %123 = (DynamicPPL.contextual_isfixed)(__context__, vn#230)::Core.Const(false)
└───        goto #24 if not %123
23 ─        Core.Const(:(σ² = (DynamicPPL.getfixed_nested)(__context__, vn#230)))
└───        Core.Const(:(goto %155))
24 ┄        goto #26 if not isassumption#231::Core.Const(true)
25 ─ %128 = Core.tuple(__context__)::Core.Const((DynamicPPL.SamplingContext{DynamicPPL.SampleFromPrior, DynamicPPL.DefaultContext, Random._GLOBAL_RNG}(Random._GLOBAL_RNG(), DynamicPPL.SampleFromPrior(), DynamicPPL.DefaultContext()),))
│    %129 = (DynamicPPL.check_tilde_rhs)(dist#233::Core.PartialStruct(Truncated{Normal{Float64}, Continuous, Float64, Float64, Nothing}, Any[Core.Const(Normal{Float64}(μ=0.0, σ=0.01)), Core.Const(0.0), Core.Const(nothing), Float64, Float64, Core.Const(1.0), Float64, Float64]))::Core.PartialStruct(Truncated{Normal{Float64}, Continuous, Float64, Float64, Nothing}, Any[Core.Const(Normal{Float64}(μ=0.0, σ=0.01)), Core.Const(0.0), Core.Const(nothing), Float64, Float64, Core.Const(1.0), Float64, Float64])
│    %130 = (DynamicPPL.unwrap_right_vn)(%129, vn#230)::Core.PartialStruct(Tuple{Truncated{Normal{Float64}, Continuous, Float64, Float64, Nothing}, AbstractPPL.VarName{:σ², typeof(identity)}}, Any[Core.PartialStruct(Truncated{Normal{Float64}, Continuous, Float64, Float64, Nothing}, Any[Core.Const(Normal{Float64}(μ=0.0, σ=0.01)), Core.Const(0.0), Core.Const(nothing), Float64, Float64, Core.Const(1.0), Float64, Float64]), AbstractPPL.VarName{:σ², typeof(identity)}])
│    %131 = Core.tuple(__varinfo__@_37)::Tuple{DynamicPPL.TypedVarInfo{@NamedTuple{p::DynamicPPL.Metadata{Dict{AbstractPPL.VarName{:p, typeof(identity)}, Int64}, Vector{DistributionsAD.TuringScalMvNormal{Vector{Float64}, Float64}}, Vector{AbstractPPL.VarName{:p, typeof(identity)}}, Vector{Float64}, Vector{Set{DynamicPPL.Selector}}}, σ²::DynamicPPL.Metadata{Dict{AbstractPPL.VarName{:σ², typeof(identity)}, Int64}, Vector{Truncated{Normal{Float64}, Continuous, Float64, Float64, Nothing}}, Vector{AbstractPPL.VarName{:σ², typeof(identity)}}, Vector{Float64}, Vector{Set{DynamicPPL.Selector}}}}, Float64}}
│    %132 = Core._apply_iterate(Base.iterate, DynamicPPL.tilde_assume!!, %128, %130, %131)::Tuple{Float64, DynamicPPL.TypedVarInfo{@NamedTuple{p::DynamicPPL.Metadata{Dict{AbstractPPL.VarName{:p, typeof(identity)}, Int64}, Vector{DistributionsAD.TuringScalMvNormal{Vector{Float64}, Float64}}, Vector{AbstractPPL.VarName{:p, typeof(identity)}}, Vector{Float64}, Vector{Set{DynamicPPL.Selector}}}, σ²::DynamicPPL.Metadata{Dict{AbstractPPL.VarName{:σ², typeof(identity)}, Int64}, Vector{Truncated{Normal{Float64}, Continuous, Float64, Float64, Nothing}}, Vector{AbstractPPL.VarName{:σ², typeof(identity)}}, Vector{Float64}, Vector{Set{DynamicPPL.Selector}}}}, Float64}}
│    %133 = Base.indexed_iterate(%132, 1)::Core.PartialStruct(Tuple{Float64, Int64}, Any[Float64, Core.Const(2)])
│           (value#234 = Core.getfield(%133, 1))
│           (@_11 = Core.getfield(%133, 2))
│    %136 = Base.indexed_iterate(%132, 2, @_11::Core.Const(2))::Core.PartialStruct(Tuple{DynamicPPL.TypedVarInfo{@NamedTuple{p::DynamicPPL.Metadata{Dict{AbstractPPL.VarName{:p, typeof(identity)}, Int64}, Vector{DistributionsAD.TuringScalMvNormal{Vector{Float64}, Float64}}, Vector{AbstractPPL.VarName{:p, typeof(identity)}}, Vector{Float64}, Vector{Set{DynamicPPL.Selector}}}, σ²::DynamicPPL.Metadata{Dict{AbstractPPL.VarName{:σ², typeof(identity)}, Int64}, Vector{Truncated{Normal{Float64}, Continuous, Float64, Float64, Nothing}}, Vector{AbstractPPL.VarName{:σ², typeof(identity)}}, Vector{Float64}, Vector{Set{DynamicPPL.Selector}}}}, Float64}, Int64}, Any[DynamicPPL.TypedVarInfo{@NamedTuple{p::DynamicPPL.Metadata{Dict{AbstractPPL.VarName{:p, typeof(identity)}, Int64}, Vector{DistributionsAD.TuringScalMvNormal{Vector{Float64}, Float64}}, Vector{AbstractPPL.VarName{:p, typeof(identity)}}, Vector{Float64}, Vector{Set{DynamicPPL.Selector}}}, σ²::DynamicPPL.Metadata{Dict{AbstractPPL.VarName{:σ², typeof(identity)}, Int64}, Vector{Truncated{Normal{Float64}, Continuous, Float64, Float64, Nothing}}, Vector{AbstractPPL.VarName{:σ², typeof(identity)}}, Vector{Float64}, Vector{Set{DynamicPPL.Selector}}}}, Float64}, Core.Const(3)])
│           (__varinfo__@_37 = Core.getfield(%136, 1))
│           (σ² = value#234)
│           value#234
└───        goto #27
26 ─        Core.Const(:((DynamicPPL.inargnames)(vn#230, __model__)))
│           Core.Const(:(!%141))
│           Core.Const(:(goto %145 if not %142))
│           Core.Const(:(σ² = (DynamicPPL.getconditioned_nested)(__context__, vn#230)))
│           Core.Const(:((DynamicPPL.check_tilde_rhs)(dist#233)))
│           Core.Const(:(σ²))
│           Core.Const(:(vn#230))
│           Core.Const(:((DynamicPPL.tilde_observe!!)(__context__, %145, %146, %147, __varinfo__@_37)))
│           Core.Const(:(Base.indexed_iterate(%148, 1)))
│           Core.Const(:(value#232 = Core.getfield(%149, 1)))
│           Core.Const(:(@_10 = Core.getfield(%149, 2)))
│           Core.Const(:(Base.indexed_iterate(%148, 2, @_10)))
│           Core.Const(:(__varinfo__@_37 = Core.getfield(%152, 1)))
└───        Core.Const(:(value#232))
27 ┄        (#11 = %new(Main.:(var"#11#12"), p@_29))
│    %156 = #11::var"#11#12"
│    %157 = (:dims,)::Core.Const((:dims,))
│    %158 = Core.apply_type(Core.NamedTuple, %157)::Core.Const(NamedTuple{(:dims,)})
│    %159 = Core.tuple(2)::Core.Const((2,))
│    %160 = (%158)(%159)::Core.Const((dims = 2,))
│    %161 = Core.kwcall(%160, Main.eachslice, mat)::Core.PartialStruct(ColumnSlices{Matrix{Float64}, Tuple{Base.OneTo{Int64}}, SubArray{Float64, 1, Matrix{Float64}, Tuple{Base.Slice{Base.OneTo{Int64}}, Int64}, true}}, Any[Matrix{Float64}, Core.Const((Colon(), 1)), Tuple{Base.OneTo{Int64}}])
│    %162 = Base.Generator(%156, %161)::Core.PartialStruct(Base.Generator{ColumnSlices{Matrix{Float64}, Tuple{Base.OneTo{Int64}}, SubArray{Float64, 1, Matrix{Float64}, Tuple{Base.Slice{Base.OneTo{Int64}}, Int64}, true}}, var"#11#12"}, Any[var"#11#12", Core.PartialStruct(ColumnSlices{Matrix{Float64}, Tuple{Base.OneTo{Int64}}, SubArray{Float64, 1, Matrix{Float64}, Tuple{Base.Slice{Base.OneTo{Int64}}, Int64}, true}}, Any[Matrix{Float64}, Core.Const((Colon(), 1)), Tuple{Base.OneTo{Int64}}])])
│           (p2 = Base.collect(%162))
│    %164 = p2::Vector
│    %165 = (σ² * Main.I)::UniformScaling{Float64}
│           (dist#238 = Main.MvNormal(%164, %165))
│    %167 = Core.apply_type(AbstractPPL.VarName, :x)::Core.Const(AbstractPPL.VarName{:x})
│    %168 = (%167)()::Core.Const(x)
│           (vn#235 = (DynamicPPL.resolve_varnames)(%168, dist#238))
│    %170 = (DynamicPPL.contextual_isassumption)(__context__, vn#235)::Core.Const(true)
└───        goto #34 if not %170
28 ─ %172 = (DynamicPPL.inargnames)(vn#235, __model__)::Any
│    %173 = !%172::Any
└───        goto #30 if not %173
29 ─        goto #31
30 ─ %176 = (DynamicPPL.inmissings)(vn#235, __model__)::Any
└───        goto #32 if not %176
31 ┄        (@_44 = true)
└───        goto #33
32 ─        (@_44 = x@_38::Vector{Float64} === Main.missing)
33 ┄        (@_43 = @_44)
└───        goto #35
34 ─        Core.Const(:(@_43 = false))
35 ┄        (isassumption#236 = @_43)
│    %185 = (DynamicPPL.contextual_isfixed)(__context__, vn#235)::Core.Const(false)
└───        goto #37 if not %185
36 ─        Core.Const(:((DynamicPPL.getfixed_nested)(__context__, vn#235)))
│           Core.Const(:(x@_38 = %187))
│           Core.Const(:(@_45 = %187))
└───        Core.Const(:(goto %220))
37 ┄        goto #39 if not isassumption#236
38 ─ %192 = Core.tuple(__context__)::Core.Const((DynamicPPL.SamplingContext{DynamicPPL.SampleFromPrior, DynamicPPL.DefaultContext, Random._GLOBAL_RNG}(Random._GLOBAL_RNG(), DynamicPPL.SampleFromPrior(), DynamicPPL.DefaultContext()),))
│    %193 = (DynamicPPL.check_tilde_rhs)(dist#238)::Any
│    %194 = (DynamicPPL.unwrap_right_vn)(%193, vn#235)::Tuple{Any, AbstractPPL.VarName}
│    %195 = Core.tuple(__varinfo__@_37)::Tuple{DynamicPPL.TypedVarInfo{@NamedTuple{p::DynamicPPL.Metadata{Dict{AbstractPPL.VarName{:p, typeof(identity)}, Int64}, Vector{DistributionsAD.TuringScalMvNormal{Vector{Float64}, Float64}}, Vector{AbstractPPL.VarName{:p, typeof(identity)}}, Vector{Float64}, Vector{Set{DynamicPPL.Selector}}}, σ²::DynamicPPL.Metadata{Dict{AbstractPPL.VarName{:σ², typeof(identity)}, Int64}, Vector{Truncated{Normal{Float64}, Continuous, Float64, Float64, Nothing}}, Vector{AbstractPPL.VarName{:σ², typeof(identity)}}, Vector{Float64}, Vector{Set{DynamicPPL.Selector}}}}, Float64}}
│    %196 = Core._apply_iterate(Base.iterate, DynamicPPL.tilde_assume!!, %192, %194, %195)::Tuple{Any, DynamicPPL.TypedVarInfo{@NamedTuple{p::DynamicPPL.Metadata{Dict{AbstractPPL.VarName{:p, typeof(identity)}, Int64}, Vector{DistributionsAD.TuringScalMvNormal{Vector{Float64}, Float64}}, Vector{AbstractPPL.VarName{:p, typeof(identity)}}, Vector{Float64}, Vector{Set{DynamicPPL.Selector}}}, σ²::DynamicPPL.Metadata{Dict{AbstractPPL.VarName{:σ², typeof(identity)}, Int64}, Vector{Truncated{Normal{Float64}, Continuous, Float64, Float64, Nothing}}, Vector{AbstractPPL.VarName{:σ², typeof(identity)}}, Vector{Float64}, Vector{Set{DynamicPPL.Selector}}}}, Float64}}
│    %197 = Base.indexed_iterate(%196, 1)::Core.PartialStruct(Tuple{Any, Int64}, Any[Any, Core.Const(2)])
│           (value#239 = Core.getfield(%197, 1))
│           (@_8 = Core.getfield(%197, 2))
│    %200 = Base.indexed_iterate(%196, 2, @_8::Core.Const(2))::Core.PartialStruct(Tuple{DynamicPPL.TypedVarInfo{@NamedTuple{p::DynamicPPL.Metadata{Dict{AbstractPPL.VarName{:p, typeof(identity)}, Int64}, Vector{DistributionsAD.TuringScalMvNormal{Vector{Float64}, Float64}}, Vector{AbstractPPL.VarName{:p, typeof(identity)}}, Vector{Float64}, Vector{Set{DynamicPPL.Selector}}}, σ²::DynamicPPL.Metadata{Dict{AbstractPPL.VarName{:σ², typeof(identity)}, Int64}, Vector{Truncated{Normal{Float64}, Continuous, Float64, Float64, Nothing}}, Vector{AbstractPPL.VarName{:σ², typeof(identity)}}, Vector{Float64}, Vector{Set{DynamicPPL.Selector}}}}, Float64}, Int64}, Any[DynamicPPL.TypedVarInfo{@NamedTuple{p::DynamicPPL.Metadata{Dict{AbstractPPL.VarName{:p, typeof(identity)}, Int64}, Vector{DistributionsAD.TuringScalMvNormal{Vector{Float64}, Float64}}, Vector{AbstractPPL.VarName{:p, typeof(identity)}}, Vector{Float64}, Vector{Set{DynamicPPL.Selector}}}, σ²::DynamicPPL.Metadata{Dict{AbstractPPL.VarName{:σ², typeof(identity)}, Int64}, Vector{Truncated{Normal{Float64}, Continuous, Float64, Float64, Nothing}}, Vector{AbstractPPL.VarName{:σ², typeof(identity)}}, Vector{Float64}, Vector{Set{DynamicPPL.Selector}}}}, Float64}, Core.Const(3)])
│           (__varinfo__@_37 = Core.getfield(%200, 1))
│           (x@_38 = value#239)
│           (@_46 = value#239)
└───        goto #42
39 ─ %205 = (DynamicPPL.inargnames)(vn#235, __model__)::Any
│    %206 = !%205::Any
└───        goto #41 if not %206
40 ─        (x@_38 = (DynamicPPL.getconditioned_nested)(__context__, vn#235))
41 ┄ %209 = (DynamicPPL.check_tilde_rhs)(dist#238)::Any
│    %210 = x@_38::Vector{Float64}
│    %211 = vn#235::AbstractPPL.VarName
│    %212 = (DynamicPPL.tilde_observe!!)(__context__, %209, %210, %211, __varinfo__@_37)::Tuple{Vector{Float64}, DynamicPPL.TypedVarInfo{@NamedTuple{p::DynamicPPL.Metadata{Dict{AbstractPPL.VarName{:p, typeof(identity)}, Int64}, Vector{DistributionsAD.TuringScalMvNormal{Vector{Float64}, Float64}}, Vector{AbstractPPL.VarName{:p, typeof(identity)}}, Vector{Float64}, Vector{Set{DynamicPPL.Selector}}}, σ²::DynamicPPL.Metadata{Dict{AbstractPPL.VarName{:σ², typeof(identity)}, Int64}, Vector{Truncated{Normal{Float64}, Continuous, Float64, Float64, Nothing}}, Vector{AbstractPPL.VarName{:σ², typeof(identity)}}, Vector{Float64}, Vector{Set{DynamicPPL.Selector}}}}, Float64}}
│    %213 = Base.indexed_iterate(%212, 1)::Core.PartialStruct(Tuple{Vector{Float64}, Int64}, Any[Vector{Float64}, Core.Const(2)])
│           (value#237 = Core.getfield(%213, 1))
│           (@_7 = Core.getfield(%213, 2))
│    %216 = Base.indexed_iterate(%212, 2, @_7::Core.Const(2))::Core.PartialStruct(Tuple{DynamicPPL.TypedVarInfo{@NamedTuple{p::DynamicPPL.Metadata{Dict{AbstractPPL.VarName{:p, typeof(identity)}, Int64}, Vector{DistributionsAD.TuringScalMvNormal{Vector{Float64}, Float64}}, Vector{AbstractPPL.VarName{:p, typeof(identity)}}, Vector{Float64}, Vector{Set{DynamicPPL.Selector}}}, σ²::DynamicPPL.Metadata{Dict{AbstractPPL.VarName{:σ², typeof(identity)}, Int64}, Vector{Truncated{Normal{Float64}, Continuous, Float64, Float64, Nothing}}, Vector{AbstractPPL.VarName{:σ², typeof(identity)}}, Vector{Float64}, Vector{Set{DynamicPPL.Selector}}}}, Float64}, Int64}, Any[DynamicPPL.TypedVarInfo{@NamedTuple{p::DynamicPPL.Metadata{Dict{AbstractPPL.VarName{:p, typeof(identity)}, Int64}, Vector{DistributionsAD.TuringScalMvNormal{Vector{Float64}, Float64}}, Vector{AbstractPPL.VarName{:p, typeof(identity)}}, Vector{Float64}, Vector{Set{DynamicPPL.Selector}}}, σ²::DynamicPPL.Metadata{Dict{AbstractPPL.VarName{:σ², typeof(identity)}, Int64}, Vector{Truncated{Normal{Float64}, Continuous, Float64, Float64, Nothing}}, Vector{AbstractPPL.VarName{:σ², typeof(identity)}}, Vector{Float64}, Vector{Set{DynamicPPL.Selector}}}}, Float64}, Core.Const(3)])
│           (__varinfo__@_37 = Core.getfield(%216, 1))
└───        (@_46 = value#237)
42 ┄        (@_45 = @_46)
│           (retval#240 = @_45)
│    %221 = Core.tuple(retval#240, __varinfo__@_37)::Tuple{Any, DynamicPPL.TypedVarInfo{@NamedTuple{p::DynamicPPL.Metadata{Dict{AbstractPPL.VarName{:p, typeof(identity)}, Int64}, Vector{DistributionsAD.TuringScalMvNormal{Vector{Float64}, Float64}}, Vector{AbstractPPL.VarName{:p, typeof(identity)}}, Vector{Float64}, Vector{Set{DynamicPPL.Selector}}}, σ²::DynamicPPL.Metadata{Dict{AbstractPPL.VarName{:σ², typeof(identity)}, Int64}, Vector{Truncated{Normal{Float64}, Continuous, Float64, Float64, Nothing}}, Vector{AbstractPPL.VarName{:σ², typeof(identity)}}, Vector{Float64}, Vector{Set{DynamicPPL.Selector}}}}, Float64}}
└───        return %221

I’ve tested both implementations and find that they are themself type stable:

p0 = randn(3)

dots_1(p, A) = [dot(p, aa) for aa ∈ eachslice(A, dims=2)]

function dots_2(p, A) 
    dots_2 = zeros(4)
    for ix ∈ axes(A, 2)
        dots_2[ix] = dot(p, A[:, ix])
    end
    dots_2
end

dots_1(p0, A) ≈ dots_2(p0, A)

@code_warntype dots_1(p0, A)
@code_warntype dots_2(p0, A)

Does anyone have an idea where this type-instability originates from?