Skip to content

VNT Part 2: Solving the issue of block elements #1180

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

Open
mhauru wants to merge 17 commits into
base: mhauru/vnt-for-fastldf
Choose a base branch
from
Open

VNT Part 2: Solving the issue of block elements #1180

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
17 commits
Select commit Hold shift + click to select a range
35c3e20
ArrayLikeBlock WIP
mhauru Dec 11, 2025
4253e9b
ArrayLikeBlock WIP2
mhauru Dec 12, 2025
5cb3916
Improve type stability of ArrayLikeBlock stuff
mhauru Dec 12, 2025
a96bb44
Test more invariants
mhauru Dec 15, 2025
a8014e6
Actually run VNT tests
mhauru Dec 15, 2025
cfc6041
Implement show for ArrayLikeBlock
mhauru Dec 15, 2025
e198fbb
Change keys on VNT to return an array
mhauru Dec 15, 2025
b77b0af
Fix keys and some tests for PartialArray
mhauru Dec 15, 2025
633e920
Improve type stability
mhauru Dec 15, 2025
222334a
Fix keys for PartialArray
mhauru Dec 15, 2025
d22face
More ArrayLikeBlock tests
mhauru Dec 15, 2025
4cb49e1
Add docstrings
mhauru Dec 15, 2025
420a6b2
Remove redundant code, improve documentation
mhauru Dec 15, 2025
d2f58d7
Merge branch 'mhauru/vnt-for-fastldf' into mhauru/arraylikeblock
mhauru Dec 15, 2025
ce9da19
Add Base.size(::RangeAndLinked)
mhauru Dec 15, 2025
4eb33e9
Fix issues with RangeAndLinked and VNT
mhauru Dec 15, 2025
51b399a
Write more design doc for ArrayLikeBlocks
mhauru Dec 15, 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
23 changes: 23 additions & 0 deletions docs/src/internals/varnamedtuple.md
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -144,6 +144,29 @@ You can also set the elements with `vnt = setindex!!(vnt, @varname(a[1]), 3.0)`,
At this point you can not set any new values in that array that would be outside of its range, with something like `vnt = setindex!!(vnt, @varname(a[5]), 5.0)`.
The philosophy here is that once a `Base.Array` has been attached to a `VarName`, that takes precedence, and a `PartialArray` is only used as a fallback when we are told to store a value for `@varname(a[i])` without having any previous knowledge about what `@varname(a)` is.

## Non-Array blocks with `IndexLens`es

The above is all that is needed for setting regular scalar values.
However, in DynamicPPL we also have a particular need for something slightly odd:
We sometimes need to do calls like `setindex!!(vnt, @varname(a[1:5]), val)` on a `val` that is _not_ an `AbstractArray`, or even iterable at all.
Normally this would error: As a scalar value with size `()`, `val` is the wrong size to be set with `@varname(a[1:5])`, which clearly wants something with size `(5,)`.
However, we want to allow this even if `val` is not an iterable, if it is some object for which `size` is well-defined, and `size(val) == (5,)`.
In DynamicPPL this comes up when storing e.g. the priors of a model, where a random variable like `@varname(a[1:5])` may be associated with a prior that is a 5-dimensional distribution.

Internally, a `PartialArray` is just a regular `Array` with a mask saying which elements have been set.
Hence we can't store `val` directly in the same `PartialArray`:
We need it to take up a sub-block of the array, in our example case a sub-block of length 5.
To this end, internally, `PartialArray` uses a wrapper type called `ArrayLikeWrapper`, that stores `val` together with the indices that are being used to set it.
The `PartialArray` has all its corresponding elements, in our example elements 1, 2, 3, 4, and, 5, point to the same wrapper object.

While such blocks can be stored using a wrapper like this, some care must be taken in indexing into these blocks.
For instance, after setting a block with `setindex!!(vnt, @varname(a[1:5]), val)`, we can't `getindex(vnt, @varname(a[1]))`, since we can't return "the first element of five in `val`", because `val` may not be indexable in any way.
Similarly, if next we set `setindex!!(vnt, @varname(a[1]), some_other_value)`, that should invalidate/delete the elements `@varname(a[2:5])`, since the block only makes sense as a whole.
Because of these reasons, setting and getting blocks of well-defined size like this is allowed with `VarNamedTuple`s, but _only by always using the full range_.
For instance, if `setindex!!(vnt, @varname(a[1:5]), val)` has been set, then the only valid `getindex` key to access `val` is `@varname(a[1:5])`;
Not `@varname(a[1:10])`, nor `@varname(a[3])`, nor for anything else that overlaps with `@varname(a[1:5])`.
`haskey` likewise only returns true for `@varname(a[1:5])`, and `keys(vnt)` only has that as an element.

## Limitations

This design has a several of benefits, for performance and generality, but it also has limitations:
Expand Down
10 changes: 4 additions & 6 deletions ext/DynamicPPLMarginalLogDensitiesExt.jl
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,6 +1,6 @@
module DynamicPPLMarginalLogDensitiesExt

using DynamicPPL: DynamicPPL, LogDensityProblems, VarName
using DynamicPPL: DynamicPPL, LogDensityProblems, VarName, RangeAndLinked
using MarginalLogDensities: MarginalLogDensities

# A thin wrapper to adapt a DynamicPPL.LogDensityFunction to the interface expected by
Expand Down Expand Up @@ -105,11 +105,9 @@ function DynamicPPL.marginalize(
ldf = DynamicPPL.LogDensityFunction(model, getlogprob, varinfo)
# Determine the indices for the variables to marginalise out.
varindices = mapreduce(vcat, marginalized_varnames) do vn
if DynamicPPL.getoptic(vn) === identity
ldf._iden_varname_ranges[DynamicPPL.getsym(vn)].range
else
ldf._varname_ranges[vn].range
end
# The type assertion helps in cases where the model is type unstable and thus
# `varname_ranges` may have an abstract element type.
(ldf._varname_ranges[vn]::RangeAndLinked).range
end
mld = MarginalLogDensities.MarginalLogDensity(
LogDensityFunctionWrapper(ldf, varinfo),
Expand Down
13 changes: 11 additions & 2 deletions src/contexts/init.jl
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -206,13 +206,17 @@ an unlinked value.

$(TYPEDFIELDS)
"""
struct RangeAndLinked
struct RangeAndLinked{T<:Tuple}
# indices that the variable corresponds to in the vectorised parameter
range::UnitRange{Int}
# whether it's linked
is_linked::Bool
# original size of the variable before vectorisation
original_size::T
Copy link
Member Author

@mhauru mhauru Dec 15, 2025

Choose a reason for hiding this comment

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

It's a bit unfortunate to have to create this field, because now you may end up with abstractly typed RangeAndLinked, which obscures the fact that the two fields you really care about, namely range and is_linked, are still concrete. The reason this is needed is that VNT needs to know how much "space" in a PartialArray an instance of RangeAndLinked takes. An alternative to this could be something like giving setindex!!(::VarNamedTuple, ...) an extra kwarg of something like ignore_size_checks.

Copy link
Member

@penelopeysm penelopeysm Dec 18, 2025 •
edited
Loading

Choose a reason for hiding this comment

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

yes, either that -- or you could just pass the size itself as a kwarg and if the kwarg is not passed, it can fall back to calling vnt_size on the value being set. not sure if that is bad for type stability as i've been bitten by julia kwargs many times but from an interface perspective, i think that is most flexible as it allows callers to specifically override the auto-size determination mechanism. my suspicion is type stability should be ok though since it will become Core.kwcall with NamedTuple{size::Tuple{Int,Int}} (or whatever the size is) which is all concretely typed.

when constructing the VNT of RangeAndLinked (inside the LDF setup code) it should be trivial to determine the size separately from the range + linked status and pass it into setindex. right now we would just need to read from the Metadata dist field, otherwise if it's an accumulator then the info is passed through from accumulate_assume!!. after the VNT has been constructed we would only ever read from the VNT, so there would no longer be any need to know the size and it can be dropped from this struct

Copy link
Member Author

@mhauru mhauru Dec 18, 2025

Choose a reason for hiding this comment

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

I agree that it makes this extra unfortunate that this is only needed when setting the value, by getindex-time it's not needed. I'm not necessarily opposed to something like the kwarg solution, though I'm not a big fan of how it's different from BangBang.setindex!! for everything else. I simultaneously hold these two conflicting opinions:

  • If some library defines a function f that operates for many types, and it uses the (semantically) same arguments for almost all of them, then if to define your own method for f you need to change the arguments, that means you shouldn't define the thing your defining as a method of f, but rather as a different function.
  • setindex!! is clearly the right function to define a method for for VNT, and anything else would be an unnecessary convolution of the interface that makes VNT less intuitive and more laborious to use.

I don't think it's an unreasonable request to say that if you want to set a value of type X into an array as a block, so that it covers some range like [1:3, 4:17, 100:101], then type X must have some suitable notion of size or otherwise it's a no-go. You could say that the notion of size something you provide ephemerally at setindex!! time, but it feels a bit hacky, compared to saying that X actually "owns" its notion of size.

Copy link
Member

@penelopeysm penelopeysm Dec 18, 2025 •
edited
Loading

Choose a reason for hiding this comment

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

That makes sense. I think, the purist part of me (which is quite a large part of me...) much prefers your first argument over your second.
Why not both, though? We could have BangBang.setindex!! for 99% of use cases, and something like DynamicPPL.setindex!!_sized that allows you to pass the size (or skip the check and infer from the varname; I'm ok with either). The implementation would be mostly shared

BangBang.setindex!!(..., val) = DynamicPPL.setindex!!_sized(..., val, vnt_size(val))

so code duplication would not be too much; and presumably one can explain in the docstring that the latter should only be used when you the programmer promises that it's a sensible operation (maybe call it unsafe or anything, but it's not really unsafe since there's no real danger in setting a wrong size, as long as you are always consistent with the indices used for getting and setting, which the code enforces).
I guess it is a bit more complexity; personally I would be OK with making that tradeoff in return for the satisfaction of knowing that we have done Good Programming Practices.

Copy link
Member

@penelopeysm penelopeysm Dec 19, 2025 •
edited
Loading

Choose a reason for hiding this comment

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

Another way of framing it is that we're going to have to introduce a hack somewhere to make this work, and I think I prefer the hack of controlling the size (or equivalently skipping the check and inferring it from the varname key; come to think of it, I'm starting to prefer that option) vs. the hack of bundling additional fields into a data structure that doesn't really need it. I don't think that RangeAndLinked intrinsically needs the size: if we were using a Dict{VarName} then there's no need for the size (although of course that comes with other drawbacks). So the way I see it, the reason for needing a hack is VNT, and it seems unfair to force RangeAndLinked to adjust itself for someone else.

end

Base.size(ral::RangeAndLinked) = ral.original_size
Copy link
Member

@penelopeysm penelopeysm Dec 18, 2025 •
edited
Loading

Choose a reason for hiding this comment

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

Not looking at the whole pr, too late, I've looked -- but wanted to comment on this line -- If one were to look at this line without all the VNT stuff context, I think it could easily look illogical -- firstly there is no broader purpose for defining Base.size on RangeAndLinked (conceptually RangeAndLinked is not a container or AbstractArray, and the range is not meant to be iterated over); secondly even if there was, this is not an unambiguously 'correct' definition of Base.size as opposed to say size(ral.range).

It appears to me that this definition only exists to satisfy the VNT interface, and for this reason I think it is better to not rely on Base.size but instead introduce a separate interface function -- call it vnt_size or anything -- that is specifically used for setting blocks in VNTs.

Using a separate function also avoids issues with structs imported from other libraries where Base.size is either not defined (and defining it would be type piracy); or is defined in a way that is not what we want it to be for VNT (we can't override the definition). As a concrete example, Base.size(::ProductNamedTupleDistribution) does not exist:

julia> size(product_distribution((; a = Normal())))
ERROR: MethodError: no method matching size(::Distributions.ProductNamedTupleDistribution{(:a,), Tuple{Normal{Float64}}, Continuous, Float64})

We could of course still have a fallback definition of

vnt_size(x) = size(x)

to avoid redefining size on everything.

Copy link
Member Author

@mhauru mhauru Dec 18, 2025

Choose a reason for hiding this comment

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

I think this is a very good idea. Will implement, not sure if I'll do it now or in January.


"""
VectorWithRanges{Tlink}(
varname_ranges::VarNamedTuple,
Expand Down Expand Up @@ -247,7 +251,12 @@ struct VectorWithRanges{Tlink,VNT<:VarNamedTuple,T<:AbstractVector{<:Real}}
end

function _get_range_and_linked(vr::VectorWithRanges, vn::VarName)
return vr.varname_ranges[vn]
# The type assertion does nothing if VectorWithRanges has concrete element types, as is
# the case for all type stable models. However, if the model is not type stable,
# vr.varname_ranges[vn] may infer to have type `Any`. In this case it is helpful to
# assert that it is a RangeAndLinked, because even though it remains non-concrete,
# it'll allow the compiler to infer the types of `range` and `is_linked`.
return vr.varname_ranges[vn]::RangeAndLinked
Copy link
Member Author

@mhauru mhauru Dec 15, 2025

Choose a reason for hiding this comment

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

Without this assertion, this model:

@model function demo_one_variable_multiple_constraints(
    ::Type{TV}=Vector{Float64}
) where {TV}
    x = TV(undef, 5)
    x[1] ~ Normal()
    x[2] ~ InverseGamma(2, 3)
    x[3] ~ truncated(Normal(), -5, 20)
    x[4:5] ~ Dirichlet([1.0, 2.0])
    return (x=x,)
end

fails the test that checks that the return type of LogDensityProblems.logdensity(ldf, x) can be inferred. It infers as Any. This happens because the VNT gets a PartialArray where three elements are RangeAndLinked{Tuple{}} and one is a ArrayLikeBlock{RangeAndLinked{Tuple{Int}}, Tuple{Int}}. That makes the element type Any, and _get_range_and_linked's return type infers as Any, and all is lost.

Now, I'm fine saying that a model that mixes x[i] and x[j:l] for the same x is pretty type unstable (note that the priors have to mix univariate and multivariate). However, I don't think that should be licence for even the logdensity return type to infer as Any, and it didn't before this PR. Hence this extra assertion, to have the type instability not be so egregious.

Note that this problem is independent of the above comment of introducing a type parameter to RangeAndLinked: Even without the type parameter, the wrapping in ArrayLikeBlock would force the PartialArray to have element type Any.

Copy link
Member

@penelopeysm penelopeysm Dec 18, 2025 •
edited
Loading

Choose a reason for hiding this comment

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

Note that this problem is independent of the above comment of introducing a type parameter to RangeAndLinked: Even without the type parameter, the wrapping in ArrayLikeBlock would force the PartialArray to have element type Any

I think in principle we should be able to do better (but maybe not in practice). Once you remove the Tuple from RangeAndLinked, the definition of _getindex on a PartialArray{Union{T,ArrayLikeBlock{T}}} suggests that it can only ever return a T or error, and should thus in theory be type stable as long as T is concrete. Of course there are two problems with this. One is assuming that Julia's type inference is capable of performing this analysis. The other is assuming that the eltype will be Union{T,ArrayLikeBlock{T}} rather than Any, which I think will depend on the semantics of things like typejoin that are beyond our control.

I'm inclined to say that if Julia is capable of inferring a return type of T from _getindex(::PartialArray{Union{T,ArrayLikeBlock{T}}}), then it may or may not become worthwhile to use our own version of typejoin that special-cases Union{T,ArrayLikeBlock{T}}, a bit like how nothing / missing is special-cased right now. But that can be investigated at another time, happy to keep the assertion right now.

Copy link
Member Author

@mhauru mhauru Dec 18, 2025

Choose a reason for hiding this comment

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

This, too, I think is an excellent idea. The problem is indeed typejoin, but I haven't thought about special casing ArrayLikeBlock like that.

end
function init(
::Random.AbstractRNG,
Expand Down
10 changes: 8 additions & 2 deletions src/logdensityfunction.jl
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -330,7 +330,10 @@ function get_ranges_and_linked_metadata(md::Metadata, start_offset::Int)
for (vn, idx) in md.idcs
is_linked = md.is_transformed[idx]
range = md.ranges[idx] .+ (start_offset - 1)
all_ranges = BangBang.setindex!!(all_ranges, RangeAndLinked(range, is_linked), vn)
orig_size = varnamesize(vn)
all_ranges = BangBang.setindex!!(
all_ranges, RangeAndLinked(range, is_linked, orig_size), vn
)
offset += length(range)
end
return all_ranges, offset
Expand All @@ -341,7 +344,10 @@ function get_ranges_and_linked_metadata(vnv::VarNamedVector, start_offset::Int)
for (vn, idx) in vnv.varname_to_index
is_linked = vnv.is_unconstrained[idx]
range = vnv.ranges[idx] .+ (start_offset - 1)
all_ranges = BangBang.setindex!!(all_ranges, RangeAndLinked(range, is_linked), vn)
orig_size = varnamesize(vn)
all_ranges = BangBang.setindex!!(
all_ranges, RangeAndLinked(range, is_linked, orig_size), vn
)
offset += length(range)
end
return all_ranges, offset
Expand Down
25 changes: 25 additions & 0 deletions src/varname.jl
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -41,3 +41,28 @@ Possibly existing indices of `varname` are neglected.
) where {s,missings,_F,_a,_T}
return s in missings
end

# TODO(mhauru) This should probably be Base.size(::VarName) in AbstractPPL.
"""
varnamesize(vn::VarName)

Return the size of the object referenced by this VarName.

```jldoctest
julia> varnamesize(@varname(a))
()

julia> varnamesize(@varname(b[1:3, 2]))
(3,)

julia> varnamesize(@varname(c.d[4].e[3, 2:5, 2, 1:4, 1]))
(4, 4)
"""
function varnamesize(vn::VarName)
l = AbstractPPL._last(vn.optic)
if l isa Accessors.IndexLens
return reduce((x, y) -> tuple(x..., y...), map(size, l.indices))
else
return ()
end
end
Loading