FlexiChains by default by penelopeysm · Pull Request #44 · rsenne/ParallelMCMC.jl

penelopeysm · 2026-06-20T23:00:15Z

Some questions worth pondering:

FlexiChains.FlexiChain{Symbol} is a massive faff. When I started writing FlexiChains it was very much with the Turing use case in mind, which is why FlexiChain{VarName} has a convenient alias VNChain. I think it probably makes sense to also make an alias for FlexiChain{Symbol}.
Add SymChain alias penelopeysm/FlexiChains.jl#248
Presumably we want to update the entire test suite to use FlexiChains instead. The question would then be what is the minimal subset of MCMCChains tests that are worth retaining separately?

penelopeysm

@rsenne Thought I'd get your opinions on a couple of points before going further.

rsenne · 2026-06-21T01:00:06Z

Thanks for this!

Here are my thoughts:

Given MCMCChains is no longer the default chains library of Turing, and soon to be this package, I'm curious if we even need to support it. The cases where someone needs it could be for ArViz, MCMCDiagnosticTools, etc.? And in any event, FlexiChains already has MCMCChains conversions yes? So practically if a user really wanted to they could just convert to a MCMCChains.Chains type via FC?
If we do drop MCMCChains support we could retype the function from Type{Tchn} to ::Type{FlexiChain{Symbol}}. This would still dispatch on both methods, but since ParallelMCMCs is more specific for the first four args and tied on the last so I think it would remove the ambiguity.

Also re your questions to ponder:

I think it probably makes sense to also make an alias

I agree

minimal subset of MCMCChains tests

If you agree with my above thoughts, none! :)

penelopeysm · 2026-06-21T10:35:27Z

You're totally correct, there's no real need to support MCMCChains. That makes life a lot easier!

penelopeysm · 2026-06-21T19:07:03Z

Okay, I got all the tests to pass locally. Apart from switching over to FlexiChains, I made a couple of design calls (which I am hoping is positive 🙂), which I'll discuss in the comments

Copilot

Pull request overview

This PR migrates ParallelMCMC’s chain output and tests from MCMCChains.Chains to FlexiChains, making FlexiChains chain types (e.g. VNChain / SymChain) the primary output target across samplers and extensions.

Changes:

Replaced MCMCChains with FlexiChains in core module dependencies and in the AbstractMCMC integration (bundle_samples, mcmcsample paths).
Updated Turing/DynamicPPL and interface tests to assert on FlexiChains APIs (parameters, extras, niters, stack=true indexing).
Updated package/test dependencies to drop MCMCChains and add FlexiChains.

Reviewed changes

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

Show a summary per file

File	Description
test/test-Turing-Integration.jl	Switches Turing integration expectations from `MCMCChains` names/indexing to FlexiChains parameter/extras APIs.
test/test-GPU-Performance.jl	Removes unused `MCMCChains` test dependency.
test/test-GPU-AD-HVP.jl	Removes unused `MCMCChains` test dependency.
test/test-DEER-Turing-Logistic.jl	Updates chain type + indexing to FlexiChains (`VNChain`/`SymChain`, `stack=true`).
test/test-DEER-Interface.jl	Migrates interface tests from `names(..., :internals)`/`size` to FlexiChains `extras`/`niters`/`nchains`.
test/test-Adaptive-MALA.jl	Migrates assertions to FlexiChains chain types and extras.
test/test-AbstractMCMC-Interface.jl	Expands test coverage for symbol vs varname keys and vector-valued params under FlexiChains.
test/Project.toml	Adds `FlexiChains` and removes `MCMCChains` from the test environment deps.
src/ParallelMCMC.jl	Replaces `using MCMCChains` with `using FlexiChains` and re-exports `sample`.
src/interface.jl	Implements FlexiChains-based chain construction for all samplers; changes default param naming to a single vector-valued `:x`.
Project.toml	Adds `FlexiChains`, removes `MCMCChains`, and updates extension trigger list.
ext/LogDensityProblemsExt.jl	Updates docs to mention FlexiChains (but still contains outdated default-name text).
ext/DynamicPPLExt.jl	Switches DynamicPPL integration to FlexiChains, adds VNChain-only bundling, and overrides construction for VarName chains.

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

codecov · 2026-06-21T19:10:29Z

Codecov Report

✅ All modified and coverable lines are covered by tests.
✅ Project coverage is 82.05%. Comparing base (7c35e3d) to head (05ec062).

❌ Your project check has failed because the head coverage (82.05%) 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      #44      +/-   ##
==========================================
+ Coverage   80.06%   82.05%   +1.98%     
==========================================
  Files           8        8              
  Lines        1164     1159       -5     
==========================================
+ Hits          932      951      +19     
+ Misses        232      208      -24

☔ View full report in Codecov by Harness.
📢 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 · 2026-06-21T19:11:43Z

+    # Wrap parameter names in the format that FlexiChains expects. Ref:
+    # https://pysm.dev/FlexiChains.jl/stable/arrays/#api-fromarray
+    if param_names === nothing
+        param_names = model.param_names
+    end
+    wrapped_param_names = if param_names === nothing
+        # Wasn't defined either as `param_names` or `model.param_names`, so make some up
+        (to_parameter(:x) => (size(vals, 2),),)
+    else
+        map(param_names) do n
+            if n isa Pair
+                to_parameter(n.first) => n.second
+            elseif n isa TKey || n isa Symbol
+                to_parameter(n)
+            else
+                throw(ArgumentError("param_names must be a collection of Pairs or $TKey, got $(typeof(n))"))
+            end
+        end
+    end


The main change is to do with how param_names are handled. This is actually not really even a big change but I think it accounts for most of the real code changes in this PR (the rest is boilerplate) so I may as well explain

Previously you could pass a vector of Symbol, and then that gets stored in MCMCChains, nice and simple. Now I've 'upgraded' it to support grouping different parameters together. Let's say you have a model with dimension 3, but conceptually you really want the first parameter to be a scalar and the other two to be grouped together as a vector. FlexiChains has a constructor that lets you do that (https://pysm.dev/FlexiChains.jl/stable/arrays/#api-fromarray).

It's a faff for the user to write all that out, but this code block lets people strip away all the FlexiChains.Parameter stuff and just write something like param_names = (:x, :y => (2,)). On top of that, this also allows us to naturally 'upgrade' a Symbol chain to a VarName chain: if the user requests

chain_type=FlexiChain{VarName}, param_names=(:x, :y => (2,))

then the :x and :y will just be 'upgraded' to @varname(x) and @varname(y). (VarNames have some benefits over Symbols in that it's easier to index into part of them, say to get y[1] out.)

Finally, the default param name has been changed: instead of scalar-valued x[i] for i in 1:D, it's now just one vector-valued parameter. That makes it easier to access the entire thing as a 3D array, while still allowing access to individual elements.

Given all this, I'd suggest that the docs should probably push users towards using VNChain as the go-to chain type, and say something like 'SymChain also works if you prefer it'? I haven't yet updated the docs, but if you're in agreement then I'll do it in the final round of this PR!

Yes this all LGTM and will be a big improvement imo

penelopeysm · 2026-06-21T19:13:18Z

-)
-    N = length(samples)
-    D = model.dim
+for TKey in (Symbol, VarName)


The cost of supporting both Symbol and VarName chains is having to do this loop over both keys (to avoid ambiguities) but since it's self-contained in this file, I think it's fine. The problem I was worried about last time was more about having to have a separate but same implementation in an MCMCChainsExt.

makes sense--loops are okay w/me but see below!

penelopeysm · 2026-06-21T19:13:47Z

        @test all(a -> a == 0.0 || a == 1.0, acc)
    end

    @testset "sample() with custom param_names" begin


You can see some examples of the param_names change works in this testset.

penelopeysm · 2026-06-21T19:16:46Z

+            s = samples[i]
+            vals[i, :] .= s.x
+            internals[i, 1] = s.logp
+            internals[i, 2] = s.accepted ? 1.0 : 0.0


Last question: Is there a need to store accepted as a F64, is it just so that internals is concrete? If it's the latter, maybe we could think about making internals a NamedTuple like (logp = [logp_matrix...], accepted = [accepted_matrix...])? That would let us store it as the natural Bool. Not one for this PR though, it's a separate thing!

Good catch. Nope no reason. I like your suggestion. I'll open a new issue about it

rsenne · 2026-06-23T01:08:26Z

Okay so I was reading over the code and I had a thought (I'm mildly sleep deprived and this idea is only half-baked so feel free to say I'm not making any sense). But i was looking at the functionality you put in _construct_flexichain and this ability to rename params like this, feels generically useful?

And so I guess IIUC the FlexiChains API (see above for why i may not)--the reason we don't hook into to_nt_and_stats/to_vnt_and_stats is because we need the functionality of our current bundle_samples to handle all the parameter names etc. This function then becomes ambiguous unless we use concrete types for arg 5 (hence the loops), but I'm curious if since the functionality in _construct_flexichain feels broadly useful, I'm curious if it should maybe live upstream? I.e., what if a pattern like this existed in FlexiChains:

# helper with same functionality as _construct_flexichain
function parameter_keys(spec, ::Type{TKey}) where {TKey}
    to_parameter(vn::VarName) = Parameter(vn)
    to_parameter(s::Symbol)   = Parameter(TKey <: VarName ? VarName{s}() : s)
    map(spec) do n
        n isa Pair             ? (to_parameter(n.first) => n.second) :
        (n isa TKey || n isa Symbol) ? to_parameter(n) :
        throw(ArgumentError("param_names must be Symbols, $TKey, or name=>shape Pairs"))
    end
end

and then add a param_names kwarg

function bundle_samples(transitions, m, s, state, ::Type{FlexiChain{Symbol}};
                        param_names=nothing, stats=missing, …)
    if param_names === nothing
        #  unchanged
    else
        keys = parameter_keys(param_names, Symbol)           # Piece 1
        arr  = stack_parameter_values(transitions)           # iters × Σdims, from the hook values
        extras = stack_stats(transitions)
        return FlexiChain{Symbol}(hcat(arr, extras), (keys..., Extra.(...)...))
    end
end

Conditioned on me making a semblance of sense, the PR as is can still proceed and this can be addressed later as these changes would need to be implemented themselves. Otherwise PR is looking good thus far!

penelopeysm · 2026-06-24T14:26:25Z

the reason we don't hook into to_nt_and_stats/to_vnt_and_stats is because we need the functionality of our current bundle_samples to handle all the parameter names etc

Yes, that and also the fact that you want to generate all the samples at once. I think the current AbstractMCMC interface kind of assumes that you're generating samples iteratively. The to_nt_and_stats etc is designed to work with AbstractMCMC.step which you aren't using.

I'm curious if it should maybe live upstream?

I think maybe. I have a couple of mild objections:

I'm a bit wary about introducing more keyword arguments to bundle_samples. None of these kwargs are defined in AbstractMCMC and so if FlexiChains methods introduce their own keyword arguments it creates an undocumented interface (a problem that's described here https://invenia.github.io/blog/2020/11/06/interfacetesting/). This is not a very strong objection because I think unfortunately FlexiChains is the only 'surviving' chains type in the Julia ecosystem, so the interface is de facto defined by FlexiChains. That's a bit of a stupid situation, but well.
It doesn't actually fix the method ambiguities. So the main benefit of upstreaming it would be to allow for reuse of the pattern. But right now there's only one library doing this. I think this might potentially change in the future but I'd rather wait until there are at least two different consumers doing this before introducing this code.

Basically, let's not overengineer things yet. (In fact I feel like to_nt_and_stats is already a bit overengineered)

rsenne · 2026-06-24T14:59:53Z

you want to generate all the samples at once

ah yes great point

assumes that you're generating samples iteratively

a very reasonable assumption until this package :)

let's not overengineer things yet

makes sense/sounds good to me

rsenne · 2026-06-25T19:33:04Z

 DifferentiationInterface = "0.7.13"
 DynamicPPL = "0.40.6, 0.41"
 Enzyme = "0.13.146"
-ForwardDiff = "1"


We should still declare a compat v for this even if removed from the Ext

Oh, is it known to not work with ForwardDiff v0?

hmm i guess not known, but i haven't tested so maybe its fine

rsenne · 2026-06-25T19:40:50Z

 """
 function ParallelMCMC.DensityModel(


So right now I have ForwardDiff as the dfault for this ext because of issue #25 and ForwardDiff worked as a workaround. I failed to document this so thats my b. though I'm not sure it makes total sense to have this as a default just given there is other backends i didn't try. So at the very least we need to either re-add in the ForwardDiff as a gate for this ext or just remove this

No worries. So personally I don't think it's necessary to have ForwardDiff as an extension trigger -- for one, maybe people don't actually want to use ForwardDiff (in which case forcing them to load it is one extra dep), and secondly it reduces discoverability. After all the main package is also agnostic and doesn't load any AD backend as a hard dep so I think the extension can do the same thing.

yes i agree let's remove as hard dep and remove the default

Hmm I was thinking it's okay to keep AutoForwardDiff as the default (it would be a bit annoying as if you don't import ForwardDiff and run it it would error, but that's no different from the main package). However if I'm not mistaken, you are suggesting to leave backend unspecified in the DynamicPPL extension, and let the user pass it via the sampler instead? If that's the case I do actually think that that's better but maybe we can do it in a separate PR?

yes that was my thinking to prevent the error, but yes thats fine we can push till next PR

Okay so for now I just reverted this so that ForwardDiff is still an extension trigger (restoring the old behaviour) and we could look at it later. I haven't thought too much about it but it might be a bit of a faff as you have to re-wrap the DynamicPPL model with the adtype from the sampler hence why I thought it made sense to defer.

penelopeysm · 2026-06-27T21:46:18Z

+        function AbstractMCMC.bundle_samples(
+            ts::Vector{<:$Ttrans},
+            model::DensityModelLDF,
+            spl::$Tspl,
+            state::$Tstate,
+            chain_type::Type{SymChain};
+            kwargs...,
+        )
+            throw(ArgumentError("FlexiChains.SymChain is not supported for DynamicPPL models; please use VNChain instead."))


So I was going to add a test for this but before I do that, I realised that this might be kind of annoying behaviour, in that if you specify the wrong chain type it will do all the sampling but only error at the end. That's no worse than the current behaviour ofc (current behaviour is to do all the sampling and then have an inscrutable MethodError at the end). Just wanted to check you're okay with that first!

hmm yeah a bit annoying, but as you said better than the earlier behavior with a legible error now. I say this is good w/ me and can be revisited later if necessary. Go ahead and add some tests!

penelopeysm · 2026-06-28T23:14:16Z

I think I have addressed or deferred everything :)

rsenne · 2026-06-28T23:37:57Z

May I humbly request tests for the missed lines? Otherwise i'll get a review on this shortly

… pr/penelopeysm/20

rsenne · 2026-06-29T00:04:06Z

Everything looks good codewise -- i added a few suggested tests :)

Co-authored-by: Ryan Senne <50930199+rsenne@users.noreply.github.com>

…into pr/penelopeysm/20

rsenne · 2026-06-29T03:24:23Z

+    @eval function AbstractMCMC.bundle_samples(
+        samples::Vector{<:ParallelMALATransition},
+        model::DensityModel,
+        sampler::ParallelMALASampler,
+        state::ParallelMALAState,
+        ::Type{FlexiChains.FlexiChain{$TKey}};
+        param_names=nothing,
+        kwargs...,
+    )
+        N = length(samples)
+        D = model.dim

-    names = if param_names !== nothing
-        param_names
-    elseif model.param_names !== nothing
-        model.param_names
-    else
-        [Symbol("x[$i]") for i in 1:D]
-    end
+        internal_names = [:logp]

-    internal_names = [:logp]
+        vals = Matrix{Float64}(undef, N, D)
+        internals = Matrix{Float64}(undef, N, 1)

-    vals = Matrix{Float64}(undef, N, D)
-    internals = Matrix{Float64}(undef, N, 1)
+        for i in 1:N
+            vals[i, :] .= samples[i].x
+            internals[i, 1] = samples[i].logp
+        end

-    for i in 1:N
-        vals[i, :] .= samples[i].x
-        internals[i, 1] = samples[i].logp
+        return _construct_flexichain($TKey, vals, internals, param_names, internal_names, model)
    end
-
-    return MCMCChains.Chains(
-        hcat(vals, internals),
-        vcat(names, internal_names),
-        Dict(:parameters => names, :internals => internal_names),
-    )
 end


Note to self: An issue for another PR became apparent to me when trying to figure out what conditions this path was exercised and I came to the realization that this is only use dwhen a non-default kwarg like thinning is used.

It took me a second to realize why but basically _sample_parallel_mala_chain needs to be retooled/extended so it supports discard_initial/thinning etc as otherwise we get kicked off the fast path. I think I just never wired this up properly. This would in turn stop routing through the generic mcmcsample path.

rsenne · 2026-06-29T03:26:04Z

Thanks for all of the amazing work on this @penelopeysm! Super happy to have FlexiChains support

penelopeysm force-pushed the main branch 2 times, most recently from 28047fa to ebd73cc Compare June 20, 2026 23:24

Move MCMCChains to extension

b0cd534

penelopeysm force-pushed the main branch from ebd73cc to b0cd534 Compare June 20, 2026 23:59

penelopeysm commented Jun 21, 2026

View reviewed changes

Comment thread ext/DynamicPPLExt.jl

Comment thread src/interface.jl Outdated

Move MCMCChains to weakdep

801aa3a

penelopeysm added 5 commits June 21, 2026 19:02

Remove MCMCChains; improve default param names

bab1b19

Tests

e9c1371

Fix tests

0fefe34

typo

966bb63

no need ForwardDiff

02634b6

penelopeysm marked this pull request as ready for review June 21, 2026 19:03

Copilot AI review requested due to automatic review settings June 21, 2026 19:03

Copilot started reviewing on behalf of penelopeysm June 21, 2026 19:04 View session

Copilot AI reviewed Jun 21, 2026

View reviewed changes

Comment thread ext/DynamicPPLExt.jl

Comment thread ext/DynamicPPLExt.jl

Comment thread ext/LogDensityProblemsExt.jl Outdated

Comment thread src/interface.jl

Comment thread src/ParallelMCMC.jl

penelopeysm commented Jun 21, 2026

View reviewed changes

penelopeysm added 3 commits June 21, 2026 20:20

Update DynamicPPLExt docstring

3471cdd

Fix LDPExt docstring

9dfadb8

Fix param_names error message

414b497

penelopeysm added 4 commits June 25, 2026 18:57

Update docs

8531ba2

More words

84ae244

Document param_names

07132ae

More docs

e9e64fb

penelopeysm requested a review from rsenne June 25, 2026 18:08

Update README

0b4c98f

rsenne requested changes Jun 25, 2026

View reviewed changes

rsenne mentioned this pull request Jun 26, 2026

Make internals a NamedTuple #49

Closed

penelopeysm added 4 commits June 27, 2026 22:32

Fix kwarg in bundle_samples; add error for SymChain

35d7780

Fix import

361d70e

Remove dead _default_param_names function

0fd15db

Tighten DynamicPPL compat bound

481000e

penelopeysm commented Jun 27, 2026

View reviewed changes

penelopeysm added 2 commits June 28, 2026 00:36

Add a test for SymChain + Turing error message

7249fd2

Re-add ForwardDiff trigger

49c6de6

penelopeysm requested a review from rsenne June 28, 2026 23:14

Merge branch 'main' of https://github.com/rsenne/ParallelMCMC.jl into…

60cdb6e

… pr/penelopeysm/20

rsenne reviewed Jun 29, 2026

View reviewed changes

Comment thread test/test-Turing-Integration.jl Outdated

Comment thread test/test-AbstractMCMC-Interface.jl

penelopeysm and others added 5 commits June 29, 2026 01:43

Apply suggestions from code review

454f7c4

Co-authored-by: Ryan Senne <50930199+rsenne@users.noreply.github.com>

Add a test for model.param_names

94d81f3

Merge branch 'main' into main

cb2e861

Merge branch 'main' of https://github.com/penelopeysm/parallelmcmc.jl …

77a1908

…into pr/penelopeysm/20

add another test because i can't read my own code

05ec062

rsenne self-requested a review June 29, 2026 03:15

rsenne reviewed Jun 29, 2026

View reviewed changes

rsenne approved these changes Jun 29, 2026

View reviewed changes

rsenne merged commit 851694b into rsenne:main Jun 29, 2026
5 of 6 checks passed

This was referenced Jun 29, 2026

Default chains backend #21

Closed

Allow thinning/discard_warmup as supported Kwargs to parallel fast path #53

Open

Internals to NamedTuple #54

Merged

Conversation

penelopeysm commented Jun 20, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

penelopeysm left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

rsenne commented Jun 21, 2026

Uh oh!

penelopeysm commented Jun 21, 2026

Uh oh!

penelopeysm commented Jun 21, 2026

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Pull request overview

Reviewed changes

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

codecov Bot commented Jun 21, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Codecov Report

Uh oh!

penelopeysm Jun 21, 2026 • 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!

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 Jun 21, 2026 • 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!

Choose a reason for hiding this comment

Uh oh!

rsenne commented Jun 23, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

penelopeysm commented Jun 24, 2026

Uh oh!

rsenne commented Jun 24, 2026

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!

Choose a reason for hiding this comment

Uh oh!

penelopeysm Jun 27, 2026 • 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!

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!

Uh oh!

Uh oh!

penelopeysm commented Jun 20, 2026 •

edited

Loading

codecov Bot commented Jun 21, 2026 •

edited

Loading

penelopeysm Jun 21, 2026 •

edited

Loading

penelopeysm Jun 21, 2026 •

edited

Loading

rsenne commented Jun 23, 2026 •

edited

Loading

penelopeysm Jun 27, 2026 •

edited

Loading