Skip to content

Improvements to VarNamedVector #1098

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 12 commits into from
Oct 29, 2025
Merged

Improvements to VarNamedVector #1098

Show file tree
Hide file tree
Changes from 1 commit
Commits
Show all changes
12 commits
Select commit Hold shift + click to select a range
150de71
Change VNV to use Dict rather than OrderedDict
mhauru Oct 28, 2025
30ac1d0
Change concretisation from map(identity, x) to a comprehension
mhauru Oct 28, 2025
4ae0c6d
Improve tighten_types!! and loosen_types!!
mhauru Oct 28, 2025
4c8b006
Fix use of set_transformed!!
mhauru Oct 28, 2025
c8b0b88
Fix push!! for VarInfos
mhauru Oct 28, 2025
1f7152b
Change the default element types in VNV to be Union{}
mhauru Oct 28, 2025
61c96b0
In untyped_vector_varinfo, don't rely on Metadata
mhauru Oct 28, 2025
2a10be9
Code style
mhauru Oct 28, 2025
bb83d93
Run formatter
mhauru Oct 29, 2025
99a7d32
In VNV, use typejoin rather than promote_type
mhauru Oct 29, 2025
513edc5
Merge branch 'mhauru/varnamedvector-speed' into mhauru/vnv-speed2
mhauru Oct 29, 2025
d30eca8
Bump patch version to 0.38.4
mhauru Oct 29, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 1 addition & 1 deletion docs/src/api.md
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -414,7 +414,7 @@ DynamicPPL.reset!
DynamicPPL.update!
DynamicPPL.insert!
DynamicPPL.loosen_types!!
DynamicPPL.tighten_types
Comment on lines 416 to -417
Copy link
Member Author

@mhauru mhauru Oct 29, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

These two don't actually ever mutate their inputs. Hence the !! is questionable. However, they sometimes return the original object, sometimes an object that shares memory with the original object, so you should use them like you use !! functions: Always catch the return value and never assume that the return value is independent of the input.

DynamicPPL.tighten_types!!
```

```@docs
Expand Down
157 changes: 94 additions & 63 deletions src/varnamedvector.jl
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -116,24 +116,24 @@ like `setindex!` and `getindex!` rather than directly accessing `vnv.vals`.

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

julia> getindex_internal(vnv, :)
8-element Vector{Real}:
8-element Vector{Float64}:
46.0
48.0
1
2
3
4
5
6
1.0
2.0
3.0
4.0
5.0
6.0
Copy link
Member

@penelopeysm penelopeysm Oct 29, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This probably makes things faster, but is this really 'correct'? The values that were inserted with setindex!! were Ints.

Copy link
Member Author

@mhauru mhauru Oct 29, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

That's a very good question, and I was wondering as well. The question comes down to whether to use this

julia> promote_type(Int, Float64)
Float64

like we do now, or this

julia> Base.promote_typejoin(Int, Float64)
Real

For an untyped VNV, converting Ints to Float64s probably makes sense, because it makes any model that mixes the two types a lot faster. However, if we move to always using typed VarInfos, maybe we wouldn't need that, assuming that people don't mix Ints and Floats in a single vector. From a correctness point of view I agree that Int would be better, although you could argue that we should respect Julia's existing choice that promote_type(Int, Float64) === Float64.

Copy link
Member

@penelopeysm penelopeysm Oct 29, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

So this is the 'classic' situation where over-tightening eltype makes a difference:

using DynamicPPL, Distributions, Random

struct RandDie <: Distributions.DiscreteUnivariateDistribution end

Base.rand(rng::Random.AbstractRNG, ::RandDie) = rand(rng, 1:6)
Distributions.logpdf(::RandDie, x::Int) = (1 <= x <= 6) ? log(1/6) : -Inf

@model function roll_die()
    die_roll ~ RandDie()
    x ~ Normal()
end
model = roll_die()

# Works
md = last(DynamicPPL.init!!(model, VarInfo(DynamicPPL.Metadata()), InitFromPrior()))
DynamicPPL.evaluate!!(model, md)

# Fails
vnv = last(DynamicPPL.init!!(model, VarInfo(DynamicPPL.VarNamedVector()), InitFromPrior()))
DynamicPPL.evaluate!!(model, vnv)

#=
ERROR: MethodError: no method matching logpdf(::RandDie, ::Float64)
The function `logpdf` exists, but no method is defined for this combination of argument types.

Closest candidates are:
  logpdf(::Chernoff, ::Real)
   @ Distributions ~/.julia/packages/Distributions/psM3H/src/univariate/continuous/chernoff.jl:155
  logpdf(::Hypergeometric, ::Real)
   @ Distributions ~/.julia/packages/Distributions/psM3H/src/univariates.jl:645
  logpdf(::RandDie, ::Int64)
   @ Main REPL[14]:1
  ...
=#

Copy link
Member

@penelopeysm penelopeysm Oct 29, 2025 •
edited
Loading

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

(See also https://pysm.dev/FlexiChains.jl/stable/why/#Heterogeneous-parameter-types which describes the same problem when MCMCChains over-concretises types -- I checked all existing Distributions subtypes of DiscreteUnivariateDistribution and all of them define a method for logpdf(dist, ::Real), so this wouldn't happen e.g. with Poisson, but in principle it can cause issues)

Copy link
Member Author

@mhauru mhauru Oct 29, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Okay, that looks like a bug. Good spot. I'll change to promote_typejoin.

Copy link
Member Author

@mhauru mhauru Oct 29, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

A note for the future if this comes up again: The difference between promote_typejoin and typejoin is that the former treats Nothing and Missing as exceptions, and returns things like

julia> Base.promote_typejoin(Missing, Float64)
Union{Missing, Float64}

rather than

julia> typejoin(Missing, Float64)
Any

I think Unions might be desirable, but it probably doesn't matter much, and promote_typejoin isn't exported, so typejoin it is.

Copy link
Member Author

@mhauru mhauru Oct 29, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Done

```
"""
struct VarNamedVector{
K<:VarName,V,TVN<:AbstractVector{K},TVal<:AbstractVector{V},TTrans<:AbstractVector
K<:VarName,V,T,KVec<:AbstractVector{K},VVec<:AbstractVector{V},TVec<:AbstractVector{T}
}
"""
mapping from a `VarName` to its integer index in `varnames`, `ranges` and `transforms`
Expand All @@ -143,7 +143,7 @@ struct VarNamedVector{
"""
vector of `VarNames` for the variables, where `varnames[varname_to_index[vn]] == vn`
"""
varnames::TVN # AbstractVector{<:VarName}
varnames::KVec

"""
vector of index ranges in `vals` corresponding to `varnames`; each `VarName` `vn` has
Expand All @@ -156,14 +156,14 @@ struct VarNamedVector{
vector of values of all variables; the value(s) of `vn` is/are
`vals[ranges[varname_to_index[vn]]]`
"""
vals::TVal # AbstractVector{<:Real}
vals::VVec

"""
vector of transformations, so that `transforms[varname_to_index[vn]]` is a callable
that transforms the value of `vn` back to its original space, undoing any linking and
vectorisation
"""
transforms::TTrans
transforms::TVec

"""
vector of booleans indicating whether a variable has been explicitly transformed to
Expand All @@ -186,14 +186,14 @@ struct VarNamedVector{

function VarNamedVector(
varname_to_index,
varnames::TVN,
varnames::KVec,
ranges,
vals::TVal,
transforms::TTrans,
vals::VVec,
transforms::TVec,
is_unconstrained=fill!(BitVector(undef, length(varnames)), 0),
num_inactive=Dict{Int,Int}();
check_consistency::Bool=CHECK_CONSISTENCY_DEFAULT,
) where {K,V,TVN<:AbstractVector{K},TVal<:AbstractVector{V},TTrans<:AbstractVector}
) where {K,V,T,KVec<:AbstractVector{K},VVec<:AbstractVector{V},TVec<:AbstractVector{T}}
if check_consistency
if length(varnames) != length(ranges) ||
length(varnames) != length(transforms) ||
Expand Down Expand Up @@ -257,7 +257,7 @@ struct VarNamedVector{
# tiny bit of thought.
end

return new{K,V,TVN,TVal,TTrans}(
return new{K,V,T,KVec,VVec,TVec}(
varname_to_index,
varnames,
ranges,
Expand All @@ -269,9 +269,9 @@ struct VarNamedVector{
end
end

function VarNamedVector{K,V}() where {K,V}
function VarNamedVector{K,V,T}() where {K,V,T}
return VarNamedVector(
OrderedDict{K,Int}(), K[], UnitRange{Int}[], V[], Any[]; check_consistency=false
Dict{K,Int}(), K[], UnitRange{Int}[], V[], T[]; check_consistency=false
)
end

Expand All @@ -280,7 +280,7 @@ end
# 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}()
VarNamedVector() = VarNamedVector{VarName,Real,An}()
function VarNamedVector(xs::Pair...; check_consistency=CHECK_CONSISTENCY_DEFAULT)
return VarNamedVector(OrderedDict(xs...); check_consistency=check_consistency)
end
Expand Down Expand Up @@ -345,6 +345,12 @@ function ==(vnv_left::VarNamedVector, vnv_right::VarNamedVector)
vnv_left.num_inactive == vnv_right.num_inactive
end

function is_concretely_typed(vnv::VarNamedVector)
return isconcretetype(eltype(vnv.varnames)) &&
isconcretetype(eltype(vnv.vals)) &&
isconcretetype(eltype(vnv.transforms))
end

getidx(vnv::VarNamedVector, vn::VarName) = vnv.varname_to_index[vn]

getrange(vnv::VarNamedVector, idx::Int) = vnv.ranges[idx]
Expand Down Expand Up @@ -562,7 +568,7 @@ to be the default vectorisation transform. This undoes any possible linking.
```jldoctest varnamedvector-reset
julia> using DynamicPPL: VarNamedVector, @varname, reset!

julia> vnv = VarNamedVector();
julia> vnv = VarNamedVector{VarName,Any,Any}();

julia> vnv[@varname(x)] = reshape(1:9, (3, 3));

Expand Down Expand Up @@ -810,7 +816,7 @@ end
# with every ! call replaced with a !! call.

"""
loosen_types!!(vnv::VarNamedVector{K,V,TVN,TVal,TTrans}, ::Type{KNew}, ::Type{TransNew})
loosen_types!!(vnv::VarNamedVector, ::Type{KNew}, ::Type{VNew}, ::Type{TNew})

Loosen the types of `vnv` to allow varname type `KNew` and transformation type `TransNew`.

Expand All @@ -821,7 +827,7 @@ transformations of type `TransNew` can be pushed to it. Some of the underlying s
shared between `vnv` and the return value, and thus mutating one may affect the other.

# See also
[`tighten_types`](@ref)
[`tighten_types!!`](@ref)

# Examples

Expand All @@ -836,7 +842,9 @@ julia> setindex_internal!(vnv, collect(1:4), @varname(y), y_trans)
ERROR: MethodError: Cannot `convert` an object of type
[...]

julia> vnv_loose = DynamicPPL.loosen_types!!(vnv, typeof(@varname(y)), typeof(y_trans));
julia> vnv_loose = DynamicPPL.loosen_types!!(
vnv, typeof(@varname(y)), Float64, typeof(y_trans)
);

julia> setindex_internal!(vnv_loose, collect(1:4), @varname(y), y_trans)

Expand All @@ -847,40 +855,57 @@ julia> vnv_loose[@varname(y)]
```
"""
function loosen_types!!(
Copy link
Member Author

@mhauru mhauru Oct 29, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I've checked and both this and tighten_types!! are compiled to no-ops if the types are already loose enough/as tight as possible.

I should probably write a test to check this. Will need to think about how to do that.

vnv::VarNamedVector, ::Type{KNew}, ::Type{TransNew}
) where {KNew,TransNew}
vnv::VarNamedVector, ::Type{KNew}, ::Type{VNew}, ::Type{TNew}
) where {KNew,VNew,TNew}
K = eltype(vnv.varnames)
Trans = eltype(vnv.transforms)
if KNew <: K && TransNew <: Trans
V = eltype(vnv.vals)
T = eltype(vnv.transforms)
if KNew <: K && VNew <: V && TNew <: T
return vnv
else
vn_type = promote_type(K, KNew)
transform_type = promote_type(Trans, TransNew)
return VarNamedVector(
Dict{vn_type,Int}(vnv.varname_to_index),
Vector{vn_type}(vnv.varnames),
vnv.ranges,
vnv.vals,
Vector{transform_type}(vnv.transforms),
vnv.is_unconstrained,
vnv.num_inactive;
check_consistency=false,
)
val_type = promote_type(V, VNew)
transform_type = promote_type(T, TNew)
# This function would work the same way if the first if statement a few lines above
# was skipped, and we only checked for the below condition. However, the first one
# is constant propagated away at compile time (at least on Julia v1.11.7), whereas
# this one isn't. Hence we keep both for performance.
return if vn_type == K && val_type == V && transform_type == T
vnv
elseif isempty(vnv)
VarNamedVector(vn_type[], val_type[], transform_type[])
else
# TODO(mhauru) We allow a `vnv` to have any AbstractVector type as its vals, but
# then here always revert to Vector.
VarNamedVector(
Dict{vn_type,Int}(vnv.varname_to_index),
Vector{vn_type}(vnv.varnames),
vnv.ranges,
Vector{val_type}(vnv.vals),
Vector{transform_type}(vnv.transforms),
vnv.is_unconstrained,
vnv.num_inactive;
check_consistency=false,
)
end
end
end

"""
tighten_types(vnv::VarNamedVector)
tighten_types!!(vnv::VarNamedVector)

Return a copy of `vnv` with the most concrete types possible.
Return a `VarNamedVector` like `vnv` with the most concrete types possible.

This function either returns `vnv` itself or new `VarNamedVector` with the same values in
it, but with the element types of various containers made as concrete as possible.

For instance, if `vnv` has its vector of transforms have eltype `Any`, but all the
transforms are actually identity transformations, this function will return a new
`VarNamedVector` with the transforms vector having eltype `typeof(identity)`.

This is a lot like the reverse of [`loosen_types!!`](@ref), but with two notable
differences: Unlike `loosen_types!!`, this function does not mutate `vnv`; it also changes
not only the key and transform eltypes, but also the values eltype.
This is a lot like the reverse of [`loosen_types!!`](@ref). Like with `loosen_types!!`, the
return value may share some of its underlying storage with `vnv`, and thus mutating one may
affect the other.

# See also
[`loosen_types!!`](@ref)
Expand All @@ -890,9 +915,9 @@ not only the key and transform eltypes, but also the values eltype.
```jldoctest varnamedvector-tighten-types
julia> using DynamicPPL: VarNamedVector, @varname, loosen_types!!, setindex_internal!

julia> vnv = VarNamedVector();
julia> vnv = VarNamedVector(@varname(x) => Real[23], @varname(y) => randn(2,2));

julia> setindex!(vnv, [23], @varname(x))
julia> vnv = delete!(vnv, @varname(y));

julia> eltype(vnv)
Real
Expand All @@ -901,7 +926,7 @@ julia> vnv.transforms
1-element Vector{Any}:
identity (generic function with 1 method)

julia> vnv_tight = DynamicPPL.tighten_types(vnv);
julia> vnv_tight = DynamicPPL.tighten_types!!(vnv);

julia> eltype(vnv_tight) == Int
true
Expand All @@ -911,17 +936,24 @@ julia> vnv_tight.transforms
identity (generic function with 1 method)
```
"""
function tighten_types(vnv::VarNamedVector)
return VarNamedVector(
Dict(vnv.varname_to_index...),
map(identity, vnv.varnames),
copy(vnv.ranges),
map(identity, vnv.vals),
map(identity, vnv.transforms),
copy(vnv.is_unconstrained),
copy(vnv.num_inactive);
check_consistency=false,
)
function tighten_types!!(vnv::VarNamedVector)
return if is_concretely_typed(vnv)
# There can not be anything to tighten, so short-circuit.
vnv
elseif isempty(vnv)
VarNamedVector()
else
VarNamedVector(
Dict(vnv.varname_to_index...),
[x for x in vnv.varnames],
vnv.ranges,
[x for x in vnv.vals],
[x for x in vnv.transforms],
vnv.is_unconstrained,
vnv.num_inactive;
check_consistency=false,
)
end
end

function BangBang.setindex!!(vnv::VarNamedVector, val, vn::VarName)
Expand Down Expand Up @@ -977,14 +1009,14 @@ function insert_internal!!(vnv::VarNamedVector, val, vn::VarName, transform=noth
if transform === nothing
transform = identity
end
vnv = loosen_types!!(vnv, typeof(vn), typeof(transform))
vnv = loosen_types!!(vnv, typeof(vn), eltype(val), typeof(transform))
insert_internal!(vnv, val, vn, transform)
return vnv
end

function update_internal!!(vnv::VarNamedVector, val, vn::VarName, transform=nothing)
transform_resolved = transform === nothing ? gettransform(vnv, vn) : transform
vnv = loosen_types!!(vnv, typeof(vn), typeof(transform_resolved))
vnv = loosen_types!!(vnv, typeof(vn), eltype(val), typeof(transform_resolved))
update_internal!(vnv, val, vn, transform)
return vnv
end
Expand Down Expand Up @@ -1219,7 +1251,6 @@ julia> subset(vnv, [@varname(x[2])]) == VarNamedVector(@varname(x[2]) => [2.0])
true
"""
function subset(vnv::VarNamedVector, vns_given::AbstractVector{<:VarName})
# NOTE: This does not specialize types when possible.
vnv_new = similar(vnv)
# Return early if possible.
isempty(vnv) && return vnv_new
Expand All @@ -1231,7 +1262,7 @@ function subset(vnv::VarNamedVector, vns_given::AbstractVector{<:VarName})
end
end

return vnv_new
return tighten_types!!(vnv_new)
end

"""
Expand Down Expand Up @@ -1430,7 +1461,7 @@ true
"""
function group_by_symbol(vnv::VarNamedVector)
symbols = unique(map(getsym, vnv.varnames))
nt_vals = map(s -> tighten_types(subset(vnv, [VarName{s}()])), symbols)
nt_vals = map(s -> tighten_types!!(subset(vnv, [VarName{s}()])), symbols)
return OrderedDict(zip(symbols, nt_vals))
end

Expand Down
6 changes: 3 additions & 3 deletions test/varnamedvector.jl
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -151,7 +151,7 @@ end
@test eltype(vnv) == Real

# Empty with types.
vnv = DynamicPPL.VarNamedVector{VarName,Float64}()
vnv = DynamicPPL.VarNamedVector{VarName,Float64,typeof(identity)}()
@test isempty(vnv)
@test eltype(vnv) == Float64
end
Expand Down Expand Up @@ -369,13 +369,13 @@ end
# Explicitly setting the transformation.
increment(x) = x .+ 10
vnv = deepcopy(vnv_base)
vnv = DynamicPPL.loosen_types!!(vnv, typeof(vn_left), typeof(increment))
vnv = DynamicPPL.loosen_types!!(vnv, typeof(vn_left), eltype(vnv), typeof(increment))
DynamicPPL.setindex_internal!(
vnv, to_vec_left(val_left .+ 100), vn_left, increment
)
@test vnv[vn_left] == to_vec_left(val_left .+ 110)

vnv = DynamicPPL.loosen_types!!(vnv, typeof(vn_right), typeof(increment))
vnv = DynamicPPL.loosen_types!!(vnv, typeof(vn_right), eltype(vnv), typeof(increment))
DynamicPPL.setindex_internal!(
vnv, to_vec_right(val_right .+ 100), vn_right, increment
)
Expand Down