Fix DynamicPPL / MCMCChains methods #1076

penelopeysm · 2025-10-15T18:28:40Z

This PR:

Reimplements the pointwise_logdensities family of functions using InitContext. NOTE They now return a MCMCChains.Chains by default rather than an OrderedDict, see Return an AbstractChains object from appliying pointwise_logdensities to Abstractchains? #688 for the motivation -- although OrderedDict output is still supported (I couldn't in good conscience allow it to return a lossy MCMCChains.Chains without leaving this option available)
Removes a handful of setval! methods that were only used for the old implementation (along with corresponding tests)
Moves the pointwise_logdensities family, as well as logjoint, loglikelihood and logprior, to MCMCChainsExt where they rightfully belong.
Fixes some overly broad logjoint, ... methods in src/simple_varinfo.jl

Closes #688

Note, while working on this I discovered more problems.

The entirety of src/model_utils.jl is really meant for MCMCChains and should also be banished to the extension (or, in my opinion, removed; but then again I think that MCMCChains should just be removed, so...). The functions there are problematic in the same way that pointwise_... were (#1072) but at least they're only for internal use, so there is no reason for them to be overloaded by other types, so unlikely to cause problems in the future.

There are also issues with the current implementation of logjoint etc. on MCMCChains.Chains, they will only work if the set of varnames is static (i.e. no varying-size / varying-names).

Personally I think these remaining issues are comparatively unimportant and can be dealt with another time (or, indeed, never dealt with).

penelopeysm · 2025-10-15T18:37:09Z

src/pointwise_logdensities.jl

-        # T is the element type of the vectors that are the values of `acc.logps`. Usually
-        # it's LogProbType.
-        T = eltype(last(fieldtypes(eltype(acc.logps))))
-        # Note that in only accumulating LogPrior, we effectively ignore logjac
-        # (since we want to return log densities that don't depend on the
-        # linking status of the VarInfo).
-        subacc = accumulate_assume!!(LogPriorAccumulator{T}(), val, logjac, vn, right)
-        push!(acc, vn, subacc.logp)
+        acc.logps[vn] = logpdf(right, val)


I don't know if there was some point in going the circuitous route. Just calling logpdf seems much easier

I suspect this may have been leftover from an earlier iteration of the code, where the whichlogprob thing was a subaccumulator rather than a Symbol.

penelopeysm · 2025-10-15T18:38:38Z

test/pointwise_logdensities.jl

-    # We'll just test one, since `pointwise_logdensities(::Model, ::AbstractVarInfo)` is tested extensively,
-    # and this is what is used to implement `pointwise_logdensities(::Model, ::Chains)`. This test suite is just
-    # to ensure that we don't accidentally break the version on `Chains`.


this comment is no longer appropriate since the varinfo method is not used to implement the chains method

penelopeysm · 2025-10-15T18:39:04Z

test/varinfo.jl

        @test typed_vi[vn_y] == 2.0
    end

-    @testset "setval!" begin


these tests are no longer relevant since the corresponding methods are removed

codecov · 2025-10-15T18:44:40Z

Codecov Report

❌ Patch coverage is 95.91837% with 2 lines in your changes missing coverage. Please review.
✅ Project coverage is 80.98%. Comparing base (9bd8f16) to head (7688065).
⚠️ Report is 1 commits behind head on breaking.

Files with missing lines	Patch %	Lines
ext/DynamicPPLMCMCChainsExt.jl	94.87%	2 Missing ⚠️

Additional details and impacted files

@@             Coverage Diff              @@
##           breaking    #1076      +/-   ##
============================================
- Coverage     82.40%   80.98%   -1.43%     
============================================
  Files            42       42              
  Lines          3791     3765      -26     
============================================
- Hits           3124     3049      -75     
- Misses          667      716      +49

☔ 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.

github-actions · 2025-10-15T18:48:10Z

DynamicPPL.jl documentation for PR #1076 is available at:
https://TuringLang.github.io/DynamicPPL.jl/previews/PR1076/

ext/DynamicPPLMCMCChainsExt.jl

mhauru · 2025-10-17T09:46:41Z

ext/DynamicPPLMCMCChainsExt.jl

+julia> # construct a chain of samples using MCMCChains
+       chain = Chains(rand(10, 2, 3), [:s, :m]);
+
+julia> loglikelihood(demo_model([1., 2.]), chain);


Is it intentional that this doctest doesn't check the output at all?

Frankly, I'm not sure. I just copy-pasted it. Very happy to change it now.

Added the output! Also changed the rand() to some fixed values to avoid making it too fragile.

src/pointwise_logdensities.jl

mhauru · 2025-10-17T09:57:38Z

src/pointwise_logdensities.jl

-        # T is the element type of the vectors that are the values of `acc.logps`. Usually
-        # it's LogProbType.
-        T = eltype(last(fieldtypes(eltype(acc.logps))))
-        # Note that in only accumulating LogPrior, we effectively ignore logjac
-        # (since we want to return log densities that don't depend on the
-        # linking status of the VarInfo).
-        subacc = accumulate_assume!!(LogPriorAccumulator{T}(), val, logjac, vn, right)
-        push!(acc, vn, subacc.logp)
+        acc.logps[vn] = logpdf(right, val)


I suspect this may have been leftover from an earlier iteration of the code, where the whichlogprob thing was a subaccumulator rather than a Symbol.

mhauru · 2025-10-17T10:04:00Z

test/pointwise_logdensities.jl

-    @test !any(Base.Fix1(subsumes, @varname(x)), keys(logpriors_pointwise))
-    @test !any(vn in keys(loglikelihoods_pointwise) for vn in vns)
-    @test all(Base.Fix1(subsumes, @varname(x)), keys(loglikelihoods_pointwise))
+    @test all(Symbol(vn) in keys(logjoints_pointwise) for vn in vns)


I see the point of moving to a Chains return type, but this does feel like a regression. I can accept that, just mourning it out loud.

I totally agree 😄. I was mulling keeping the default as Dict, and make Chains output opt-in. What do you think?

(It's also less breaking.)

I do think long-term Chains feels like the right thing. We could choose to wait until MCMCChains has been improved, so that this would no longer be a regression. I'm quite ambivalent though.

If it makes a difference (not that I think it should, really), I already implemented pointwise_... for FlexiChains and they return a FlexiChain: https://github.com/penelopeysm/FlexiChains.jl/blob/66b973e0fb346f3ef6cb51a6caa70324ef9d8f70/ext/FlexiChainsDynamicPPLExt.jl#L237-L254

Co-authored-by: Markus Hauru <[email protected]>

mhauru

I'm happy to merge. Happy either way on which one is the default, so change it if you want to.

penelopeysm · 2025-10-17T16:45:05Z

Thanks! OK, Chains it is then

penelopeysm added 4 commits October 15, 2025 18:45

Reimplement pointwise_logdensities (almost completely)

7e37432

Move logjoint, logprior, ... as well

4002b08

Fix imports, etc

b444419

Remove tests that are failing (yes I learnt this from Claude)

e1489ba

penelopeysm changed the base branch from main to breaking October 15, 2025 18:28

github-actions bot assigned penelopeysm Oct 15, 2025

Changelog

2e69e0b

penelopeysm mentioned this pull request Oct 15, 2025

Return an AbstractChains object from appliying pointwise_logdensities to Abstractchains? #688

Closed

logpdf

fc393bc

penelopeysm commented Oct 15, 2025

View reviewed changes

fix docstrings

f2a83b3

penelopeysm added 3 commits October 15, 2025 19:53

allow dict output

0dec616

changelog

f41106e

fix some comments

0b1d6a6

penelopeysm marked this pull request as draft October 16, 2025 13:11

penelopeysm added 2 commits October 16, 2025 15:25

fix tests

94432eb

Fix more imports

57e099a

penelopeysm marked this pull request as ready for review October 16, 2025 15:16

penelopeysm requested a review from mhauru October 16, 2025 15:16

mhauru requested changes Oct 17, 2025

View reviewed changes

penelopeysm and others added 2 commits October 17, 2025 11:19

Remove stray n

802e38f

Co-authored-by: Markus Hauru <[email protected]>

Expand logprior, loglikelihood, and logjoint docstrings

7688065

penelopeysm requested a review from mhauru October 17, 2025 11:10

mhauru approved these changes Oct 17, 2025

View reviewed changes

penelopeysm merged commit e4c5bbd into breaking Oct 17, 2025
17 of 19 checks passed

penelopeysm deleted the py/redo-pointwise branch October 17, 2025 16:45

Fix DynamicPPL / MCMCChains methods #1076

Fix DynamicPPL / MCMCChains methods #1076

Uh oh!

Conversation

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

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!

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

Uh oh!

Codecov Report

Uh oh!

github-actions bot commented Oct 15, 2025

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!

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

Uh oh!

Choose a reason for hiding this comment

Uh oh!

mhauru left a comment

Choose a reason for hiding this comment

Uh oh!

penelopeysm commented Oct 17, 2025

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

penelopeysm commented Oct 15, 2025 •

edited

Loading

codecov bot commented Oct 15, 2025 •

edited

Loading

penelopeysm Oct 17, 2025 •

edited

Loading