Move `hasvalue` and `getvalue` from DynamicPPL; implement extra Distributions-based methods #125

penelopeysm · 2025-07-05T22:34:24Z

Julia minimum version bump

I bumped to 1.10 as I don't want to add extra code to handle extensions on pre-1.9. Most important packages in TuringLang are already using >= 1.10 anyway.

Moving functions from DynamicPPL

This PR moves hasvalue and getvalue from DynamicPPL to AbstractPPL. https://github.com/TuringLang/DynamicPPL.jl/blob/92f6eea8660be2142fa4087e5e025f37026bfa45/src/utils.jl#L763-L954

A lot of the helper functions in DynamicPPL are not actually needed because there is existing functionality in here that accomplishes much the same. I modified the implementations accordingly.

Distributions-based methods

This part is new and warrants more explanation. To begin, notice the default behaviour of hasvalue:

julia> d = Dict(@varname(x[1]) => 1.0, @varname(x[2]) => 1.0)
Dict{VarName{:x, Accessors.IndexLens{Tuple{Int64}}}, Float64} with 2 entries:
  x[1] => 1.0
  x[2] => 1.0

julia> hasvalue(d, @varname(x))
false

This makes sense, because d alone does not give us enough information to reconstruct some arbitrary variable x.

However, let's say that we know x is to be sampled from a given distribution dist. In this case, we do have enough information to determine whether x can be reconstructed. This PR therefore also implements the following methods:

julia> using Distributions, LinearAlgebra

julia> hasvalue(d, @varname(x), MvNormal(zeros(2), I))
true

julia> getvalue(d, @varname(x), MvNormal(zeros(2), I))
[1.0, 1.0]

julia> hasvalue(d, @varname(x), MvNormal(zeros(3), I))
false

The motivation for this is to (properly) fix issues where values for multivariate distributions are specified separately, see e.g., TuringLang/DynamicPPL.jl#712, see also this comment TuringLang/DynamicPPL.jl#710 (comment).

One might argue that we should force users to specify things properly, i.e., if x ~ MvNormal(zeros(2), I) then the user should condition on Dict(@varname(x) => [1.0, 1.0]) rather than Dict(@varname(x[1]) => 1.0, @varname(x[2]) => 1.0). In an ideal world I would do that, and even now, I would still advocate for making this general guideline clear in e.g. the docs.

However, there remains one specific case where this isn't enough, namely in DynamicPPL's predict(model, chain) or returned(model, chain). These methods require extracting variable values from chain, inserting them into a VarInfo, and rerunning the model with the given values. Unfortunately, chain is a lossy storage format, because array-valued variables like x are split up into x[1] and x[2] and it's not possible to recover the original shape of x.

Up until this PR, this has been handled in DynamicPPL using the setval_and_resample! and nested_setindex_maybe methods which perform some direct manipulation of VarInfos. I think these methods are slightly dangerous and can lead to subtle bugs, for example, if only part of the variable x is given, it marks the entire variable x as to be not-resampled: https://github.com/TuringLang/DynamicPPL.jl/blob/92f6eea8660be2142fa4087e5e025f37026bfa45/src/varinfo.jl#L2177-L2181

The good news, though, is that when evaluating a model, we have access to the distribution that x is supposed to be sampled from. Thus, we can determine whether enough of the x[i]'s are given to reconstruct it, which is what these new methods do. So, we can deal with this in a more principled fashion: if we can find all the indices needed to reconstruct the value of x, then we can confidently set that value; if we can't, then we don't even attempt to set any of the individual indices because hasvalue will return false.

Remaining questions:

I wonder if we can simplify the API? Note that hasvalue and getvalue have extremely similar logic, do we really need to have two functions with almost the same implementation? I've held off on attempting to do this because I'm worried about type stability, i.e. getvalue is inherently type-unstable, and maybe guarding calls to getvalue behind a call to hasvalue avoids leaking type instability into the caller function. However, I think this is reliant on the compiler being able to infer the return value of hasvalue through e.g. constant propagation?!
Not sure if this should be a minor bump. According to semver, nothing in here is breaking, hence I did patch bump. But the changes are quite large and maybe it feels more correct to do a minor bump.

TODO

This PR doesn't support ProductNamedTupleDistribution. It shouldn't be overly complicated to implement IMO. However, almost nothing else in TuringLang works with ProductNamedTupleDistribution, so I don't feel bad not implementing it.

Closes #124

This is required for the InitContext PR TuringLang/DynamicPPL.jl#967 as ParamsInit needs to use hasvalue and getvalue. Specifically, I also want to use ParamsInit to handle predict, hence the need for the Distributions-based methods.

codecov · 2025-07-05T22:35:51Z

Codecov Report

Attention: Patch coverage is 85.38462% with 19 lines in your changes missing coverage. Please review.

Project coverage is 86.28%. Comparing base (7be9556) to head (e20190c).

Files with missing lines	Patch %	Lines
ext/AbstractPPLDistributionsExt.jl	76.92%	15 Missing ⚠️
src/hasvalue.jl	93.02%	3 Missing ⚠️
src/varname.jl	95.45%	1 Missing ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##             main     #125      +/-   ##
==========================================
+ Coverage   83.56%   86.28%   +2.72%     
==========================================
  Files           2        5       +3     
  Lines         292      401     +109     
==========================================
+ Hits          244      346     +102     
- Misses         48       55       +7

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:

❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

coveralls · 2025-07-06T11:06:20Z

Pull Request Test Coverage Report for Build 16221573874

Details

111 of 130 (85.38%) changed or added relevant lines in 3 files are covered.
No unchanged relevant lines lost coverage.
Overall coverage increased (+2.7%) to 86.284%

Changes Missing Coverage	Covered Lines	Changed/Added Lines	%
src/varname.jl	21	22	95.45%
src/hasvalue.jl	40	43	93.02%
ext/AbstractPPLDistributionsExt.jl	50	65	76.92%

Totals
Change from base Build 16098285788:	2.7%
Covered Lines:	346
Relevant Lines:	401

💛 - Coveralls

github-actions · 2025-07-06T18:56:08Z

AbstractPPL.jl documentation for PR #125 is available at:
https://TuringLang.github.io/AbstractPPL.jl/previews/PR125/

penelopeysm · 2025-07-07T14:34:22Z

src/varname.jl

-Remove identity lenses from composed optics.
-"""
-_strip_identity(::Base.ComposedFunction{typeof(identity),typeof(identity)}) = identity
-function _strip_identity(o::Base.ComposedFunction{Outer,typeof(identity)}) where {Outer}
-    return _strip_identity(o.outer)
-end
-function _strip_identity(o::Base.ComposedFunction{typeof(identity),Inner}) where {Inner}
-    return _strip_identity(o.inner)
-end
-_strip_identity(o::Base.ComposedFunction) = o
-_strip_identity(o::Accessors.PropertyLens) = o
-_strip_identity(o::Accessors.IndexLens) = o
-_strip_identity(o::typeof(identity)) = o


normalise strips identities now so this function isn't needed any more

penelopeysm · 2025-07-07T14:37:11Z

src/varname.jl

+_head(o::ComposedFunction{Outer,Inner}) where {Outer,Inner} = o.inner
+_head(o::Accessors.PropertyLens) = o
+_head(o::Accessors.IndexLens) = o
+_head(::typeof(identity)) = identity


_head, _tail, _init, and _last take their names from the equivalent Haskell functions on linked lists:

λ> head [1,2,3] 1 λ> tail [1,2,3] [2,3] λ> init [1,2,3] [1,2] λ> last [1,2,3] 3 -- empty list is turned into identity in our case λ> head [1] 1 λ> tail [1] [] λ> init [1] [] λ> last [1] 1

penelopeysm · 2025-07-09T22:06:49Z

src/hasvalue.jl

+    else
+        error("getvalue: $(vn) was not found in the values provided")
+    end


I'm actually a bit on the fence about this error message. From the perspective of AbstractPPL, it makes perfect sense. But when we use this in DynamicPPL we get things like this:

julia> vi[@varname(x)] - ERROR: KeyError: key x not found + ERROR: getvalue: x was not found in the values provided

which seems to unnecessarily leak details of internal implementation.

Previously in DynamicPPL it would throw a KeyError which made sense when calling getindex. However, it doesn't make sense to throw a KeyError here because getvalue isn't indexing into a dictionary. So I'm struggling a bit.

I don't think the error message is bad. Error messages quite often reveal implementation details, I don't think that can be avoided except by using a lot of try-catch higher up (in this case in DPPL), which I think is usually far too much complexity to be worth it. My only alternative would be "$(vn) was not found in the NamedTuple provided".

Also, if we could skip the canview call, I wonder if that could have a (small but) noticeable impact on performance, due to not bounds checking twice. Or maybe the return optic(vals[sym]) call could be done with @inbounds?

Error message changed, I think I do prefer a bit not mentioning the function name (it's in the stacktrace anyway).

Re. canview: I think this ties into one of the things I wasn't fully sure about (on the main comment).

The idea is that you could use getvalue without having checked hasvalue, so in principle you do need to check bounds in both. That is very annoying not only because of the double bounds check but also because of the severe code duplication (notice that this PR could have been half the lines if we combined the two functions into one).

I don't fully know how to solve this though. The obvious solution would be for getvalue to return a sentinel value if it's not found. However I worry that that would introduce type instability.

mhauru

Happy with the code except for some minor localised comments, but confused about whether we really need to deal with LKJCholesky, and thus whether we could restrict ourselves to just caring about dimensions rather than more complicated forms of distribution outputs.

HISTORY.md

mhauru · 2025-07-11T09:11:33Z

HISTORY.md

+These functions check whether a given `VarName` has a value in the given `NamedTuple` or `AbstractDict`, and return the value if it exists.
+
+The optional `Distribution` argument allows one to reconstruct a full value from its component indices.
+For example, if `container` has `x[1]` and `x[2]`, then `hasvalue(container, @varname(x), dist)` will return true if `size(dist) == (2,)` (for example, `MvNormal(zeros(2), I)`).


Did you consider having the third argument be the dimension, rather than the distribution? I'm not sure at all that this would be better, but it would avoid a dependence on Distributions.jl for hasvalue, and I was wondering if it has merit.

We were so close, it's just Cholesky that breaks it.

There is another option actually, which is instead of taking a Distribution, take a value sampled from that Distribution. That means we would only need a dependency on LinearAlgebra rather than Distributions.

However, that would mean an additional call to rand() which, while quite minor in the grand scheme of things, I feel opposed to in principle.

src/hasvalue.jl

mhauru · 2025-07-11T09:23:20Z

src/hasvalue.jl

+    else
+        error("getvalue: $(vn) was not found in the values provided")
+    end


I don't think the error message is bad. Error messages quite often reveal implementation details, I don't think that can be avoided except by using a lot of try-catch higher up (in this case in DPPL), which I think is usually far too much complexity to be worth it. My only alternative would be "$(vn) was not found in the NamedTuple provided".

Also, if we could skip the canview call, I wonder if that could have a (small but) noticeable impact on performance, due to not bounds checking twice. Or maybe the return optic(vals[sym]) call could be done with @inbounds?

mhauru · 2025-07-11T10:08:31Z

ext/AbstractPPLDistributionsExt.jl

+function get_optics(dist::Distributions.LKJCholesky)
+    is_up = dist.uplo == 'U'
+    cartesian_indices = filter(CartesianIndices(size(dist))) do cartesian_index
+        i, j = cartesian_index.I
+        is_up ? i <= j : i >= j
+    end
+    # there is an additional layer as we need to access `.L` or `.U` before we
+    # can index into it
+    field_lens = is_up ? (Accessors.@o _.U) : (Accessors.@o _.L)
+    return map(idx -> Accessors.IndexLens(idx.I) ∘ field_lens, cartesian_indices)
+end


I'm confused by the scenario in which we need this. LKJCholesky returns objects of type Cholesky, which are not AbstractArrays and can't be indexed. Would we ever have a situation where we would have something like

getvalue(Dict(@varname(x[1,1]) => 1.0, @varname(x[1,2]) => 0.0, @varname(x[2,2]) => 1.0), @varname(x), LKJCholesky(2, 0.5))

where it should return true?

PS. After reading the tests I know realise you need things like @varname(x.U[1,1]), but I still wonder if this would ever come up.

I guess this ambiguity could come up if the user manually mis-specified without the .L or .U. However, assuming that the varnames come from MCMCChains, then the varnames will have been constructed correctly (via varname_and_value_leaves... which ALSO special-cases Cholesky https://github.com/TuringLang/DynamicPPL.jl/blob/ce7c8b1ae48624e12ebf3064b6099e3dfca8c985/src/utils.jl#L1258-L1265)

mhauru · 2025-07-11T10:21:02Z

ext/AbstractPPLDistributionsExt.jl

+) where {sym}
+    # If `vn` is present as-is, then we are good
+    AbstractPPL.hasvalue(vals, vn) && return true
+    # If not, then we need to check inside `vals` to see if a subset of
+    # `vals` is enough to reconstruct `vn`. For example, if `vals` contains
+    # `x[1]` and `x[2]`, and `dist` is `MvNormal(zeros(2), I)`, then we
+    # can reconstruct `x`. If `dist` is `MvNormal(zeros(3), I)`, then we
+    # can't.
+    # To do this, we get the size of the distribution and iterate over all
+    # possible indices. If every index can be found in `subsumed_keys`, then we
+    # can return true.
+    optics = get_optics(dist)
+    original_optic = AbstractPPL.getoptic(vn)
+    expected_vns = map(o -> VarName{sym}(o ∘ original_optic), optics)
+    if all(sub_vn -> AbstractPPL.hasvalue(vals, sub_vn), expected_vns)
+        return true
+    else
+        if error_on_incomplete &&
+            any(sub_vn -> AbstractPPL.hasvalue(vals, sub_vn), expected_vns)
+            error("hasvalue: only partial values for `$vn` found in the values provided")
+        end
+        return false
+    end


What happens, and what should happen, with the following?

hasvalue( Dict(@varname(x) => [1.0], @varname(x[2]) => 2.0), @varname(x), MvNormal(zeros(2), I), )

I think this returns true on line 163 even though the shape of @varname(x) => [1.0] doesn't match the distribution. But even if it didn't, I think it would it also return true because @varname(x[1]) is found in @varname(x) => [1.0] and @varname(x[2]) is found in @varname(x[2]) => 2.0. That doesn't quite feel right though.

(cf. comment below in main thread)

test/hasvalue.jl

mhauru · 2025-07-11T10:29:13Z

test/hasvalue.jl

+        @test hasvalue(d, @varname(x), MvNormal(zeros(1), I))
+        @test getvalue(d, @varname(x), MvNormal(zeros(1), I)) == [1.0]


It's not wrong, but it feels funny.

(cf. comment below in main thread)

penelopeysm · 2025-07-11T12:49:38Z

Hmm. I guess maybe the biggest overall takeaway is that although the implementation makes sense, there is an overarching theme of ill-defined behaviour for wrongly specified input, which I totally get, I felt the same way writing it.

The problem is that this ill-defined behaviour already used to exist in e.g. DynamicPPL's nested_setindex_maybe. You would run into such behaviour if for example you passed in a chain with the wrong variables (like if you had MvNormal(zeros(1)) and passed in a chain with x[1] and x[2], I'm not sure what it would do with x[2]).

Assuming that such ill-defined behaviour needs to exist somewhere in order for us to keep compatibility with predict(..., ::MCMCChains), how about the following compromise solution?

This is where hasvalue(..., dist) gets (or will get) used upstream (TuringLang/DynamicPPL.jl#981):

https://github.com/TuringLang/DynamicPPL.jl/blob/b55c1e17f97ae518d1d149122e1fb1055557183f/src/contexts/init.jl#L93-L110

Right now, it just unconditionally uses the dist method. Maybe we can include a flag to ParamsInit which tells us when we allow using the dist method, and restrict its use to specific cases where we absolutely require this behaviour, i.e. predict.

That would forbid, for example, people sampling and specifying initial values that don't align with the model.

Co-authored-by: Markus Hauru <markus@mhauru.org>

penelopeysm marked this pull request as draft July 5, 2025 22:34

github-actions bot assigned penelopeysm Jul 5, 2025

Base automatically changed from py/composed-assoc to main July 6, 2025 10:58

penelopeysm added 3 commits July 6, 2025 12:04

Move hasvalue and getvalue to AbstractPPL; reimplement

56b9f2c

Add hasvalue for (some) distributions

646b9d7

Bump min Julia to 1.10

07975a2

penelopeysm force-pushed the py/hasgetvalue branch from 9310e27 to 07975a2 Compare July 6, 2025 11:04

penelopeysm added 4 commits July 6, 2025 13:25

Make hasvalue and getvalue use the most specific value

332c64a

Specify getvalue semantics in docstring

049001e

Simplify logic (can rely on normalisation)

e0adba7

Add tests for composition of head/tail and init/last

9291e07

penelopeysm force-pushed the py/hasgetvalue branch from 1e3ed7c to 9291e07 Compare July 6, 2025 13:19

penelopeysm mentioned this pull request Jul 6, 2025

Setting values of triangular matrices also sets data above/below the diagonal JuliaObjects/Accessors.jl#203

Closed

penelopeysm added 2 commits July 6, 2025 19:47

Finish implementing distributions methods

a736e70

Document

5a902b0

TuringLang deleted a comment from github-actions bot Jul 6, 2025

penelopeysm added 3 commits July 6, 2025 19:56

Fix LinearAlgebra version bound

0e8d256

Try to fix documentation for extension (why is this so complicated...)

29dc922

Fix extension documentation

a998af6

penelopeysm marked this pull request as ready for review July 6, 2025 19:03

penelopeysm force-pushed the py/hasgetvalue branch from 6fe03f9 to 398b42f Compare July 6, 2025 19:40

Implement fallback {has,get}value methods for NamedTuple + Distribution

6a01588

penelopeysm force-pushed the py/hasgetvalue branch from 398b42f to 6a01588 Compare July 6, 2025 23:32

penelopeysm requested a review from mhauru July 7, 2025 10:34

penelopeysm commented Jul 7, 2025

View reviewed changes

This was referenced Jul 9, 2025

key type of SimpleVarInfo is underfined TuringLang/DynamicPPL.jl#791

Open

InitContext, part 2 - Move hasvalue and getvalue to AbstractPPL; enforce key type of AbstractDict TuringLang/DynamicPPL.jl#980

Draft

Fix wrong way round composition, add more tests

ecf0eac

penelopeysm force-pushed the py/hasgetvalue branch from 3a254b4 to ecf0eac Compare July 9, 2025 22:01

penelopeysm commented Jul 9, 2025

View reviewed changes

mhauru reviewed Jul 11, 2025

View reviewed changes

penelopeysm and others added 8 commits July 11, 2025 14:28

Update src/hasvalue.jl

8dc53d8

Co-authored-by: Markus Hauru <markus@mhauru.org>

Update HISTORY.md

a85e0e0

Co-authored-by: Markus Hauru <markus@mhauru.org>

Minor bump

7fac592

Fix test (forgot to push this...)

5fa6a53

Add extra example for getvalue

b4f69b1

Tweak error message when value not found

00520c3

Format (?!)

d25f0be

Fix doctests

e20190c

penelopeysm requested a review from mhauru July 11, 2025 15:49

		@test hasvalue(d, @varname(x), MvNormal(zeros(1), I))
		@test getvalue(d, @varname(x), MvNormal(zeros(1), I)) == [1.0]

Move hasvalue and getvalue from DynamicPPL; implement extra Distributions-based methods #125

Are you sure you want to change the base?

Move hasvalue and getvalue from DynamicPPL; implement extra Distributions-based methods #125

Uh oh!

Conversation

penelopeysm commented Jul 5, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Julia minimum version bump

Moving functions from DynamicPPL

Distributions-based methods

Remaining questions:

TODO

Uh oh!

codecov bot commented Jul 5, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Codecov Report

Uh oh!

coveralls commented Jul 6, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Pull Request Test Coverage Report for Build 16221573874

Details

💛 - Coveralls

Uh oh!

github-actions bot commented Jul 6, 2025

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

mhauru left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

penelopeysm commented Jul 11, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

Uh oh!

Move `hasvalue` and `getvalue` from DynamicPPL; implement extra Distributions-based methods #125

Move `hasvalue` and `getvalue` from DynamicPPL; implement extra Distributions-based methods #125

penelopeysm commented Jul 5, 2025 •

edited

Loading

codecov bot commented Jul 5, 2025 •

edited

Loading

coveralls commented Jul 6, 2025 •

edited

Loading

penelopeysm commented Jul 11, 2025 •

edited

Loading