`vi[:]` unnecessarily casts

[penelopeysm]

julia> using DynamicPPL, Distributions

julia> @model function g()
           x ~ Poisson(3.5)
           y ~ Normal()
       end;

julia> v = VarInfo(g());

julia> DynamicPPL.getindex_internal(v, @varname(x))
1-element Vector{Int64}:
 1

julia> v[:]
2-element Vector{Float64}:
  1.0
 -0.7525359545725246

This is was causing upstream test failures in Turing because DynamicPPL.unflatten(v, v[:]) can set float values where ints should be present. I worked around it for now by avoiding round-trip unflatten, but this still seems like something that should be fixed here.

[penelopeysm]

I remember the last time I attempted to twiddle with this code, I ran into type stability issues, so anybody who takes this on should probably be wary of that.

[mhauru]

Another downside of our need to adhere to LogDensityProblems' interface and thus vectorise a whole VarInfo into and out of a single vector.

[penelopeysm]

Making a Vector is one thing, but I wish we could avoid unnecessarily coercing to a narrower type than needed. The resulting type would of course have to be abstract:

julia> x = Real[1, 2.0]
2-element Vector{Real}:
 1
 2.0

julia> typeof(x[1])
Int64

julia> typeof(x[2])
Float64

julia> map(identity, x)
2-element Vector{Real}:
 1
 2.0

and thus for type stability issues one might prefer to have something like NamedTuple, but at least you wouldn't have this issue:

julia> using DynamicPPL, Distributions

julia> @model function g()
           x ~ Poisson(3.5)
           y ~ Normal()
       end;

julia> v = VarInfo(g());

julia> v[@varname(x)]
4

julia> vnew = DynamicPPL.unflatten(v, v[:]);

julia> vnew[@varname(x)]
4.0

If you attempt to then do DynamicPPL.evaluate!!(model, vnew), it does work, but that's only because Distributions is kind enough to make logpdf(::Poisson, ::Float64) return the 'correct' value if the float is a whole number. If it only defined logpdf(::Poisson, ::Int), then evaluate!! would error.

There are some things in Turing.RandomMeasures which only define logpdf on Ints which is where I first ran into this issue.

[penelopeysm]

Equivalently I guess the upshot of ditching LogDensityProblems is that we could potentially do this in a type-stable way.

[penelopeysm]

Another place where unflatten and vi[:] get used is in LogDensityFunction:

DynamicPPL.jl/src/logdensityfunction.jl

Lines 179 to 184 in ce7c8b1

	function logdensity_at(
	x::AbstractVector, model::Model, varinfo::AbstractVarInfo, context::AbstractContext
	)
	varinfo_new = unflatten(varinfo, x)
	return getlogp(last(evaluate!!(model, varinfo_new, context)))
	end

If you naively do LogDensityProblems.logdensity(ldf, ldf.varinfo[:]) (i.e. 'calculate the log density at the current point'), then you can get errors like this: #1002 (comment)

using DynamicPPL, StableRNGs, Distributions, ADTypes
using DynamicPPL.TestUtils.AD: run_ad, NoTest

@model function lda(K, d, w)
    V = length(unique(w))
    D = length(unique(d))
    N = length(d)
    @assert length(w) == N
    ϕ = Vector{Vector{Real}}(undef, K)
    for i in 1:K
        ϕ[i] ~ Dirichlet(ones(V) / V)
    end
    θ = Vector{Vector{Real}}(undef, D)
    for i in 1:D
        θ[i] ~ Dirichlet(ones(K) / K)
    end
    z = zeros(Int, N)
    for i in 1:N
        z[i] ~ Categorical(θ[d[i]])
        w[i] ~ Categorical(ϕ[d[i]])
    end
    return (; ϕ=ϕ, θ=θ, z=z)
end
w = [1, 2, 3, 2, 1, 1]
d = [1, 1, 1, 2, 2, 2]
model = lda(2, d, w)

# Works
vi = DynamicPPL.typed_varinfo(StableRNG(23), model)
vi = DynamicPPL.link!!(vi, model)
DynamicPPL.evaluate!!(model, vi)

#= Fails 
ERROR: BoundsError: attempt to access 4-element Vector{Float64} at index [4:5]
Stacktrace:
  [1] throw_boundserror(A::Vector{Float64}, I::Tuple{UnitRange{Int64}})
    @ Base ./essentials.jl:14
  [2] checkbounds
    @ ./abstractarray.jl:699 [inlined]
=#
p = vi[:]
vi_new = DynamicPPL.unflatten(vi, p)
DynamicPPL.evaluate!!(model, vi_new)