Map DynamicPPL samples back to original space

[penelopeysm]

Hello! So it turns out I did manage to work on this a little bit. This PR just takes care of mapping parameter values back into the 'original' space and getting back their name.

I think it's easiest to demonstrate this with the following model, which I snuck a variation of into the test suite. There are a couple of things here: firstly, the transformed Dirichlet will only have two variables, but in the chain the result should have three. Secondly, the parameter names should take their cue from the distributions: x is a length-3 vector, and y is a NamedTuple with keys :a and :b.

using Turing, DynamicPPL, MCMCChains
@model function f()
    x ~ Dirichlet(ones(3))
    y ~ product_distribution((a = Normal(), b = Normal()))
end

sample(f(), NUTS(), 2000; chain_type=MCMCChains.Chains)
# parameters        = x[1], x[2], x[3], y.a, y.b

The way this works is by reevaluating the model:

julia> ldf = LogDensityFunction(f(), getlogjoint_internal, LinkAll());

julia> ParamsWithStats(randn(4), ldf) # This reevaluates.
ParamsWithStats
 ├─ params
 │  VarNamedTuple
 │  ├─ x => [0.10010881155549031, 0.46021694655404727, 0.43967424189046245]
 │  └─ y => (a = -0.013561686835955641, b = -0.1326274207437508)
 └─ stats
    ├─ logprior = -1.1536168618908884
    ├─ loglikelihood = 0.0
    └─ logjoint = -1.1536168618908884

What this PR does is slightly different for the different samplers:

• for anything that overloads AbstractMCMC.step, we inject this reevaluation at the end of the step via a new function _postprocess_sample.
• for the combination of ParallelMALASampler + MCMCChains.Chains, because you've currently overloaded AbstractMCMC.mcmcsample directly, we inject it via a new function _construct_chains.
• sometimes we also need to overload bundle_samples directly :(. Ugly, but well.

All of these new functions are overloaded in DynamicPPLExt. However, we have to make sure to only overload them for DensityModels which ultimately wrap a DynamicPPL.LogDensityFunction. To make this possible, we have to thread through some extra structs that allow us to dispatch on this, namely the LogDensityProblemPrimal and LogDensityProblemGradient structs (previously the logp and gradlogp field would be anonymous functions which can't be dispatched on).

Overall, this means that the following works nicely (this goes via AbstractMCMC.step):

using ParallelMCMC, DynamicPPL, MCMCChains, Distributions
@model function f()
    x ~ Dirichlet(ones(3))
    y ~ product_distribution((a = Normal(), b = Normal()))
end

dmodel = DensityModel(f());
spl = AdaptiveMALASampler(0.3; n_warmup=500)

sample(dmodel, spl, 2000; chain_type=MCMCChains.Chains, discard_warmup=false)
# Samples per chain = 2000
# parameters        = x[1], x[2], x[3], y.a, y.b

sample(dmodel, spl, 2000; chain_type=MCMCChains.Chains, discard_warmup=true)
# Samples per chain = 1499
# parameters        = x[1], x[2], x[3], y.a, y.b

The DEER model in the test suite also works correctly too, however I couldn't get ParallelMALASampler to work nicely on Turing models out of the box because Enzyme's hvp errorred in various forms. I assume that the test suite works because you supply an analytical hvp...

As for FlexiChains, the fact that we overload AbstractMCMC.step means that things sometimes Just Work:

sample(dmodel, spl, 2000; chain_type=FlexiChains.VNChain, discard_warmup=false)
# ↓ iter  = 1:2000
# │ Parameters (2) ── VarName
# │  Vector{Float64}           x
# │  @NamedTuple{a::Float64,…  y

However, discard_warmup=true doesn't work: that's because we haven't overloaded bundle_samples for FlexiChains yet. I haven't put in too much effort into this yet, because I think that that should be separate -- we can at least get MCMCChains working for now.

[Copilot]

Pull request overview

This PR updates the DynamicPPL/Turing integration so samples produced in linked/unconstrained space are postprocessed back into the model’s original parameter space (including original-shaped parameters and names), primarily targeting MCMCChains.Chains outputs.

Changes:

• Introduces a postprocess_sample hook and applies it in AbstractMCMC.step implementations to allow remapping transitions after sampling steps.
• Adds dispatchable wrappers for LogDensityProblems primal/gradient callables to enable DynamicPPL-specific postprocessing.
• Extends chain construction/bundling paths (notably for ParallelMALASampler + MCMCChains) and updates tests + compat/version.

Reviewed changes

Copilot reviewed 5 out of 6 changed files in this pull request and generated 5 comments.

Show a summary per file

File	Description
src/interface.jl	Adds postprocessing hook and threads it through sampler steps; refactors ParallelMALASampler chain construction to allow specialization.
ext/LogDensityProblemsExt.jl	Wraps LDP objects in dispatchable callable structs; extends constructor to forward hvp.
ext/DynamicPPLExt.jl	Implements DynamicPPL-specific postprocessing and MCMCChains bundling/chain construction based on ParamsWithStats.
test/test-Turing-Integration.jl	Adds coverage for constrained + structured parameters mapping back to original space and names.
test/test-DEER-Turing-Logistic.jl	Adjusts Turing/DEER tests for the updated naming/mapping behavior.
Project.toml	Bumps package version and expands DynamicPPL compatibility bounds.

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

[rsenne]

Thanks so much for getting this started! I'll take a look this week as time permits :)

I couldn't get ParallelMALASampler to work nicely on Turing models out of the box because Enzyme's hvp errorred in various forms. I assume that the test suite works because you supply an analytical hvp...

by any chance could you send me an example error? I thought i had just got this fixed...might have to call in an AD expert if this ends up being that brittle. On my machine at home I was able to get it working without error so it's unfortunate to hear its broken elsewhere...

[codecov]

Codecov Report

❌ Patch coverage is 89.28571% with 3 lines in your changes missing coverage. Please review.
✅ Project coverage is 88.65%. Comparing base (b99f19d) to head (f58b89f).
⚠️ Report is 15 commits behind head on main.

Files with missing lines	Patch %	Lines
ext/DynamicPPLExt.jl	83.33%	3 Missing ⚠️

❌ Your patch status has failed because the patch coverage (89.28%) is below the target coverage (90.00%). You can increase the patch coverage or adjust the target coverage.
❌ Your project status has failed because the head coverage (88.65%) is below the target coverage (90.00%). You can increase the head coverage or adjust the target coverage.

Additional details and impacted files

@@            Coverage Diff             @@
##             main      #20      +/-   ##
==========================================
- Coverage   88.89%   88.65%   -0.25%     
==========================================
  Files           6        6              
  Lines        1045     1040       -5     
==========================================
- Hits          929      922       -7     
- Misses        116      118       +2     

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

[penelopeysm]

No rush at all :) Yes, I'll file a separate issue for Enzyme.

[rsenne]

Overall looks great, just one point and one quick question for you :)

[penelopeysm]

Sorry it took a while! This week turned out to be pretty busy for me.

DynamicPPLExt basically now has one overload mirroring each overload in src/interface.jl. I condensed some of them into an eval block, which I prefer (just by a tiny bit), but I'm happy to change that to just duplicate the code instead, if you'd rather.

I added a couple sentences to the docs.

[rsenne]

@all-contributors please add @rsenne for code,maintenance,test,ideas,review,docs

Just getting this up-to-date as i forgot to do this

[rsenne]

@all-contributors please add @penelopeysm for code,test,ideas,review,docs

[allcontributors]

@rsenne

I've put up a pull request to add @penelopeysm! 🎉

[rsenne]

LGTM. Thanks so much for your work and thanks for being the first contributor :)

[penelopeysm]

thanks for being the first contributor

Ooh, it is my honour!