`VarNamedTuple`

[penelopeysm]

A fundamental 'problem' in DynamicPPL is the choice between NamedTuple and Dict. For example, ConditionContext and FixedContext can be instantiated with both. There is also the entirety of the typed VarInfo data structure which is heavily influenced by this.

The main benefit of NamedTuple is type stability, which is accomplished through the use of macros and/or generated functions. The main drawback of NamedTuple is the inability to represent VarNames with a non-identity optic (e.g. @varname(x[1]), @varname(x.a)) correctly. The pros and cons of Dict are the opposite.

It would be very nice if we could make a new data structure that combines the best of both worlds, i.e. represents a mapping from (possibly non-identity) VarNames to arbitrary data, but can be used in a type-stable (and reasonably performant) manner.

This would help to eliminate a bunch of things that sometimes work and sometimes don't.

It might also open up possibilities for further tidying up metadata / VNV data structures to ensure that information stays in sync.

I haven't investigated this concretely yet, but @mhauru and I have been mucking around with this idea for a while now.

[yebai]

The motivations appear to be very similar to those underlying VarNamedVector. What would be the differences between VarNamedVector and VarNamedTuple?

edit: VarNamedTuple would allow us to use different types for each field, but VarNamedVector would default to vector of floats.

edit: DynamicPPL's tracing data structures are similar to JAX's pytree, so looking at pytree might provide useful ideas.

[yebai]

Some further notes:

VarNamedVector's getindex returns a vector of Real, so it won't preserve types, unlike NamedTuple. For example:

julia> vnv[@varname(x)]
2-element Vector{Real}:
 46.0
 48.0

The most recent Julia releases (version 1.8 and later) introduced support for type declarations on variables and function arguments (see the Julia documentation on Type Declarations). These declarations provide crucial type information to the compiler, enabling it to ensure type stability.

This feature opens up a potential avenue to address the type-preservation issue in VarNamedVector. By adding a type parameter to the getindex method, such as getindex(vnv, @varname(x), T), we could use typeassert(ret, T) to guarantee that the returned vector has the desired type T.

This approach suggests that it might be feasible to utilize VarNamedVector directly within typed_vector_varinfo by leveraging these type declarations and typeassert, thereby potentially eliminating the need for a NamedTuple-of-VarNamedVector structure to achieve type stability. Successfully implementing this could pave the way for a unification of typed and untyped VarInfo representations.

`typeassert` example

function foo(x, T)
  ret = x > 0 ? 1.0 : 2
  typeassert(ret, T)
end

julia> @code_warntype foo(2., Int)
MethodInstance for foo(::Float64, ::Type{Int64})
  from foo(x, T) @ Main REPL[3]:1
Arguments
  #self#::Core.Const(Main.foo)
  x::Float64
  T::Core.Const(Int64)
Locals
  ret::Union{Float64, Int64}
  @_5::Union{Float64, Int64}
Body::Int64
1 ─       Core.NewvarNode(:(ret))
│   %2  = Main.:>::Core.Const(>)
│   %3  = (%2)(x, 0)::Bool
└──       goto #3 if not %3
2 ─       (@_5 = 1.0)
└──       goto #4
3 ─       (@_5 = 2)
4 ┄ %8  = @_5::Union{Float64, Int64}
│         (ret = %8)
│   %10 = ret::Union{Float64, Int64}
│   %11 = Main.typeassert(%10, T)::Int64
└──       return %11


julia> @code_warntype foo(2., Float64)
MethodInstance for foo(::Float64, ::Type{Float64})
  from foo(x, T) @ Main REPL[3]:1
Arguments
  #self#::Core.Const(Main.foo)
  x::Float64
  T::Core.Const(Float64)
Locals
  ret::Union{Float64, Int64}
  @_5::Union{Float64, Int64}
Body::Float64
1 ─       Core.NewvarNode(:(ret))
│   %2  = Main.:>::Core.Const(>)
│   %3  = (%2)(x, 0)::Bool
└──       goto #3 if not %3
2 ─       (@_5 = 1.0)
└──       goto #4
3 ─       (@_5 = 2)
4 ┄ %8  = @_5::Union{Float64, Int64}
│         (ret = %8)
│   %10 = ret::Union{Float64, Int64}
│   %11 = Main.typeassert(%10, T)::Float64
└──       return %11

[penelopeysm]

I think I don't primarily envision this being used as the backbone of VarInfo itself - I think it is too much hassle to change the data structure of VarInfo itself. I think it's more likely to come into play with things like (in increasing order of speculativeness)

• conditioning and fixing
• setting initial values for sampling
• calculating logp at a given point ie the type of x in logjoint(model, x)
• the default output for rand(model)
• a possible return type for submodel distributions (see comments in this PR, which won't get merged by now, but still useful as a reminder of where I got to)
• it might potentially be useful as a replacement for AbstractMCMC.Transition
• in VarInfo, storing things like the istrans flag together with their associated VarName (rather than in a separate array)

Just to take the top example (since it's easy), ConditionContext right now can wrap either a NamedTuple or a Dict.

DynamicPPL.jl/src/contexts.jl

Lines 368 to 373 in 6cd439d

	struct ConditionContext{
	Values<:Union{NamedTuple,AbstractDict{<:VarName}},Ctx<:AbstractContext
	} <: AbstractContext
	values::Values
	context::Ctx
	end

I think that:

1. This is confusing.
2. ConditionContext should internally always wrap a single data structure, not one of two.
3. It's ok to retain the convenience of allowing users to condition with a NamedTuple like model | (x = 1.0, y = 2.0), but internally this should be converted to a single data structure.
4. Neither NamedTuple nor Dict are perfect data structures, hence the need for a new one.

As for VNV vs this, it's true that VNV lets you index by varnames, sometimes in a type-stable way. The thing about it returning Real is because of these lines

DynamicPPL.jl/src/varnamedvector.jl

Lines 266 to 271 in 23834de

	# TODO(mhauru) I would like for this to be VarNamedVector(Union{}, Union{}). Simlarly the
	# transform vector type above could then be Union{}[]. This would allow expanding the
	# VarName and element types only as necessary, which would help keep them concrete. However,
	# making that change here opens some other cans of worms related to how VarInfo uses
	# BangBang, that I don't want to deal with right now.
	VarNamedVector() = VarNamedVector{VarName,Real}()

Personally, the biggest difference to me is just the data structure itself:

1. VNV pretends to be a mapping of varnames to things. It does quite a good job at masquerading as that, but it's really not: it's a mapping of varnames to index ranges, which are then used to index a vector. What this means is that we can't treat a varname, plus the data associated with it, as a single unit - it creates an extra step between the key and the value. The main reason why it has these ranges is presumably to deal with linking properly. However, this complexity isn't needed for most of the purposes described above.
2. VNV fields are static (because it's a struct) and it can't be generalised to holding different types of information.

[mhauru]

I think this would be great for things like condition. E.g. with GibbsConditionContext I found myself at some point essentially replicating TypedVarInfo's "type stable at the level of leading symbol of a VarName" logic and thought "we should have a data structure for this".

For VarInfo, I originally thought this would replace exactly the NamedTuple thing in TypedVarInfo. So the metadata field of a VarInfo would be a VarNamedTuple with values of VarNamedVector. If we could get it to somehow deal nicely with index ranges too (so not just typed on the symbol of the VarName) then indeed we could consider building an AbstractVarInfo completly using it, too. I think that would be essentially SimpleVarInfo with this new data structure as a backend (SimpleVarInfo is another place where we mix Dicts and NamedTuples), with the main different to VarInfo and VNV being indeed that it handles linking differently using DynamicTransform (support for linking is quite deeply built-in in VNV).