Sketch of VarNamedTuple

[mhauru]

This PR is unlikely ever to be merged. I presume once the features here are done, they will need to be split into smaller PRs, and some of it may go into AbstractPPL. I'm putting this up just do the development in public and be able to talk about it.

[penelopeysm]

I like what I've read so far. There is a tricky thing with make_leaf, which is that IndexLenses are not necessarily disjoint (unlike PropertyLenses with different symbols). That is, we could make make_leaf("a", @o _[1])[@o _[1]] and make_leaf("b", @o _[2])[@o _[2]] work, but then we don't know how to deal with ...[@o _[1:2]] or ...[@o _[:]], which could all be conceivable operations. I don't immediately know how to fix this, but my impression is still that we need to carry information about the size / shape / type of the variables. That is we need to know what @varname(x) is before we can safely return a value for @varname(x[1:2]) or @varname(x[:]). That would I think remove the need for make_leaf, at the cost of complexity elsewhere. I think this is workable as long as IndexLenses are the innermost / tail lens.

There are probably about 5 different ways this approach could also go awry, so I'm not quite sure.

[mhauru]

I need to give more thought to the idea of passing the value of the variable to tilde_assume!!. That could solve a lot of problems, but I worry that it might open new ones. For instance, if

v = Vector{Float64}(undef, 3)
v[1] ~ Normal()

pushes all of v into __varinfo__ then do we need to keep track of the fact that v[2] and v[3] are still not valid keys? Also, __varinfo__ would need to be able to hold v without knowing what the prior distribution or bijector of some of the elements are, which is a new situation we haven't yet had to handle.

[mhauru]

I've been fixing my sketch implementation and adding some workaround hacks to Metadata, to see if I can get reasonable performance out of this approach. Here's a snippet I've been using to benchmark

module TmpTest

using Random
using DynamicPPL
using Distributions
using BenchmarkTools
using LinearAlgebra

@model function f(dim)
    x ~ Normal(0, 1)
    y ~ Normal(x, 1)

    z = Vector{Float64}(undef, dim)
    for i in 1:dim
        z[i] ~ Normal(y, 1)
    end

    data ~ MvNormal(z, I)
    return nothing
end

dim = 300
m = f(dim) | (; data=randn(dim))

Random.seed!(23)
vi_tuple = DynamicPPL.tuple_varinfo(m)

Random.seed!(23)
vi_typed = DynamicPPL.typed_varinfo(m)

# svi_nt = DynamicPPL.SimpleVarInfo(vi_typed, NamedTuple)

svi_od = DynamicPPL.SimpleVarInfo(vi_typed, OrderedDict)

vi_tuple = link(vi_tuple, m)
vi_typed = link(vi_typed, m)
# svi_nt = link(svi_nt, m)
svi_od = link(svi_od, m)

logp = getlogjoint(last(DynamicPPL.evaluate!!(m, vi_tuple)))
@show logp
logp = getlogjoint(last(DynamicPPL.evaluate!!(m, vi_typed)))
@show logp
# logp = getlogjoint(last(DynamicPPL.evaluate!!(m, svi_nt)))
# @show logp
logp = getlogjoint(last(DynamicPPL.evaluate!!(m, svi_od)))
@show logp

# println("Link times")
# @btime link(vi_tuple, m)
# @btime link(vi_typed, m)
# # @btime link(svi_nt, m)
# @btime link(svi_od, m)

println("Evaluate times")
@btime DynamicPPL.evaluate!!(m, vi_tuple)
@btime DynamicPPL.evaluate!!(m, vi_typed)
# @btime DynamicPPL.evaluate!!(m, svi_nt)
@btime DynamicPPL.evaluate!!(m, svi_od)

end

The summary is that, with these hacks, I can take a model with varnames like @varname(x) and indexed ones like @varname(z[1]), make a VarInfo that is a VNT with Metadata objects at the leaves, link it, and evaluate the model on it. For smallish models performance is wihin a factor of 2 of our current NTVarInfo type, sometimes better, sometimes worse. What I'm learning is that the extra indirection of the machinery in VNT, compared to just a single NamedTuple, seems to have little cost to it.

Number of Symbols seems to not be an issue. I just made a model with 42 distinct variables all with their own lead Symbols, resulting in a NamedTuple with 42 keys, and all seems fine with it, including type inference. As long as we don't start turning IndexLenses into Symbols we are good.

[github-actions]

[JuliaFormatter v1.0.62] _{reported by reviewdog 🐶}

Suggested change

	* People might really need IndexLenses in the middle of VarNames. The one place this comes up is submodels within a loop. I'm still inclined to keep designing without allowing for that, for now, but should keep in mind that that needs to be relaxed eventually. If it makes it easier, we can require that users explicitly tell us the size of any arrays for which this is done.
	* When storing values for nested NamedTuples, the actual variable may be a struct. Do we need to be able to reconstruct the struct from the NamedTuple? If so, how do we do that?
	* Do `Colon` indices cause any extra trouble for the leafnodes?
	- People might really need IndexLenses in the middle of VarNames. The one place this comes up is submodels within a loop. I'm still inclined to keep designing without allowing for that, for now, but should keep in mind that that needs to be relaxed eventually. If it makes it easier, we can require that users explicitly tell us the size of any arrays for which this is done.
	- When storing values for nested NamedTuples, the actual variable may be a struct. Do we need to be able to reconstruct the struct from the NamedTuple? If so, how do we do that?
	- Do `Colon` indices cause any extra trouble for the leafnodes?

[mhauru]

#1097 prompted me to wonder, if I do

vi = TupleVarInfo()
vi = setindex!!(vi, 1.0, @varname(x[1]))
getindex(vi, @varname(x))

what should the last line return?

I would like to say it should return [1.0], but what if I've only pushed a value for @varname(x[3]), rather than @varname(x[1])? Maybe [missing, missing, 1.0]?