From 1424c5e061c45760e34e53307afd4356714170d9 Mon Sep 17 00:00:00 2001 From: Ryan Senne <50930199+rsenne@users.noreply.github.com> Date: Tue, 30 Jun 2026 22:57:19 -0400 Subject: [PATCH 1/4] Start `NamedTuples` switch --- Project.toml | 2 + ext/DynamicPPLExt.jl | 7 ++-- src/ParallelMCMC.jl | 1 + src/interface.jl | 93 ++++++++++++++++++++++++++------------------ 4 files changed, 62 insertions(+), 41 deletions(-) diff --git a/Project.toml b/Project.toml index d422312..acf4ccf 100644 --- a/Project.toml +++ b/Project.toml @@ -10,6 +10,7 @@ CUDA = "052768ef-5323-5732-b1bb-66c8b64840ba" DifferentiationInterface = "a0c0ee7d-e4b9-4e03-894e-1c5f64a51d63" FlexiChains = "4a37a8b9-6e57-4b92-8664-298d46e639f7" LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e" +OrderedCollections = "bac558e1-5e72-5ebc-8fee-abe8a469f55d" Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c" Statistics = "10745b16-79ce-11e8-11f9-7d13ad32a3b2" @@ -39,6 +40,7 @@ ForwardDiff = "1" LinearAlgebra = "1" LogDensityProblems = "2" Mooncake = "0.5.26" +OrderedCollections = "1" Random = "1" Statistics = "1" Zygote = "0.7.10" diff --git a/ext/DynamicPPLExt.jl b/ext/DynamicPPLExt.jl index 7800d81..cc33ad4 100644 --- a/ext/DynamicPPLExt.jl +++ b/ext/DynamicPPLExt.jl @@ -136,13 +136,12 @@ end function ParallelMCMC._construct_flexichain( ::Type{VarName}, vals::AbstractMatrix{<:Real}, - internals::AbstractMatrix{<:Real}, + internals::NamedTuple, ::Any, - internal_names::Vector{Symbol}, model::DensityModelLDF, ) - pwss = map(zip(eachrow(vals), eachrow(internals))) do (val, internal) - stats = NamedTuple{Tuple(internal_names)}(internal) + pwss = map(enumerate(eachrow(vals))) do (i, val) + stats = map(v -> v[i], internals) DynamicPPL.ParamsWithStats(val, model.logdensity.ld, stats) end return AbstractMCMC.from_samples(VNChain, hcat(pwss)) diff --git a/src/ParallelMCMC.jl b/src/ParallelMCMC.jl index 7a9a557..4ad005e 100644 --- a/src/ParallelMCMC.jl +++ b/src/ParallelMCMC.jl @@ -4,6 +4,7 @@ using AbstractMCMC using CUDA using FlexiChains using LinearAlgebra +using OrderedCollections: OrderedDict using Random using Statistics diff --git a/src/interface.jl b/src/interface.jl index 7523297..adf24d0 100644 --- a/src/interface.jl +++ b/src/interface.jl @@ -206,19 +206,19 @@ for TKey in (Symbol, VarName) N = length(samples) D = model.dim - internal_names = [:logp, :accepted] - vals = Matrix{Float64}(undef, N, D) - internals = Matrix{Float64}(undef, N, 2) + logp = Vector{Float64}(undef, N) + accepted = Vector{Bool}(undef, N) for i in 1:N s = samples[i] vals[i, :] .= s.x - internals[i, 1] = s.logp - internals[i, 2] = s.accepted ? 1.0 : 0.0 + logp[i] = s.logp + accepted[i] = s.accepted end - return _construct_flexichain($TKey, vals, internals, param_names, internal_names, model) + internals = (logp=logp, accepted=accepted) + return _construct_flexichain($TKey, vals, internals, param_names, model) end end @@ -563,10 +563,9 @@ function _sample_parallel_mala_chain( progressname="Sampling", ) where {TKey} D = model.dim - internal_names = [:logp] vals = Matrix{Float64}(undef, N, D) - internals = Matrix{Float64}(undef, N, 1) + logp = Vector{Float64}(undef, N) progress = _parallel_mala_progress(progress, progressname) x0 = _parallel_mala_initial_x(rng, model, sampler, initial_params) @@ -586,7 +585,7 @@ function _sample_parallel_mala_chain( rows = first_row:(first_row + nkeep - 1) _copy_trajectory_rows!(vals, first_row, S, nkeep) - internals[rows, 1] .= view(logps, 1:nkeep) + logp[rows] .= view(logps, 1:nkeep) x0 = copy(view(S, :, nkeep)) nsteps += nkeep @@ -599,15 +598,15 @@ function _sample_parallel_mala_chain( end end - return _construct_flexichain(TKey, vals, internals, param_names, internal_names, model) + internals = (logp=logp,) + return _construct_flexichain(TKey, vals, internals, param_names, model) end function _construct_flexichain( ::Type{TKey}, vals::AbstractMatrix{<:Real}, - internals::AbstractMatrix{<:Real}, + internals::NamedTuple, param_names::Any, - internal_names::Vector{Symbol}, model::DensityModel, ) where {TKey} # Wrap user-supplied names in `Parameter`. This allows people to specify, e.g., @@ -616,36 +615,55 @@ function _construct_flexichain( to_parameter(vn::VarName) = FlexiChains.Parameter(vn) to_parameter(s::Symbol) = FlexiChains.Parameter(TKey <: VarName ? VarName{s}() : s) - # vals and internals are both `iters x params` - # FlexiChains expects `iters x chains x params` - arr = hcat(vals, internals) - arr = reshape(arr, size(arr, 1), 1, size(arr, 2)) + N, D = size(vals) - # Wrap parameter names in the format that FlexiChains expects. Ref: - # https://pysm.dev/FlexiChains.jl/stable/arrays/#api-fromarray + # Normalize parameter names to `key => dims` pairs (scalars use `()`), so we can + # consume the columns of `vals` uniformly below. 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),),) + (to_parameter(:x) => (D,),) 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) + to_parameter(n) => () else throw(ArgumentError("param_names must be a collection of Pairs, Symbols, or $TKey, got $(typeof(n))")) end end end - all_names = ( - wrapped_param_names..., - FlexiChains.Extra.(internal_names)... + # Build a dict of per-key arrays rather than a single stacked array. This lets each + # key keep its natural element type — e.g. `accepted`/`is_warmup` stay `Bool` instead + # of being widened to `Float64` (see issue #49). + data = OrderedDict{FlexiChains.ParameterOrExtra{<:TKey},AbstractArray}() + + offset = 1 + for (key, sz) in wrapped_param_names + ncols = prod(sz) # `prod(()) == 1`, i.e. scalar consumes one column + if sz === () + data[key] = vals[:, offset] + else + cols = offset:(offset + ncols - 1) + data[key] = [collect(reshape(view(vals, i, cols), sz)) for i in 1:N] + end + offset += ncols + end + offset - 1 == D || throw( + ArgumentError( + "param_names consume $(offset - 1) columns, but there are $D parameter columns" + ), ) - return FlexiChains.FlexiChain{TKey}(arr, all_names) + + for name in keys(internals) + data[FlexiChains.Extra(name)] = internals[name] + end + + return FlexiChains.FlexiChain{TKey}(N, 1, data) end function _sample_parallel_mala_blocks( @@ -871,17 +889,16 @@ for TKey in (Symbol, VarName) N = length(samples) D = model.dim - internal_names = [:logp] - vals = Matrix{Float64}(undef, N, D) - internals = Matrix{Float64}(undef, N, 1) + logp = Vector{Float64}(undef, N) for i in 1:N vals[i, :] .= samples[i].x - internals[i, 1] = samples[i].logp + logp[i] = samples[i].logp end - return _construct_flexichain($TKey, vals, internals, param_names, internal_names, model) + internals = (logp=logp,) + return _construct_flexichain($TKey, vals, internals, param_names, model) end end @@ -1076,20 +1093,22 @@ for TKey in (Symbol, VarName) N = length(filtered) D = model.dim - internal_names = [:logp, :accepted, :step_size, :is_warmup] - vals = Matrix{Float64}(undef, N, D) - internals = Matrix{Float64}(undef, N, 4) + logp = Vector{Float64}(undef, N) + accepted = Vector{Bool}(undef, N) + step_size = Vector{Float64}(undef, N) + is_warmup = Vector{Bool}(undef, N) for i in 1:N s = filtered[i] vals[i, :] .= s.x - internals[i, 1] = s.logp - internals[i, 2] = s.accepted ? 1.0 : 0.0 - internals[i, 3] = s.step_size - internals[i, 4] = s.is_warmup ? 1.0 : 0.0 + logp[i] = s.logp + accepted[i] = s.accepted + step_size[i] = s.step_size + is_warmup[i] = s.is_warmup end - return _construct_flexichain($TKey, vals, internals, param_names, internal_names, model) + internals = (logp=logp, accepted=accepted, step_size=step_size, is_warmup=is_warmup) + return _construct_flexichain($TKey, vals, internals, param_names, model) end end From 8016f237ace9bfc875fe3abf70ac51cd47217fc9 Mon Sep 17 00:00:00 2001 From: Ryan Senne <50930199+rsenne@users.noreply.github.com> Date: Wed, 1 Jul 2026 08:40:36 -0400 Subject: [PATCH 2/4] Add tests and remove widening by Float64 --- src/interface.jl | 26 ++++++++++------- test/test-AbstractMCMC-Interface.jl | 44 +++++++++++++++++++++++++++-- test/test-Adaptive-MALA.jl | 11 +++++--- 3 files changed, 65 insertions(+), 16 deletions(-) diff --git a/src/interface.jl b/src/interface.jl index adf24d0..e185ead 100644 --- a/src/interface.jl +++ b/src/interface.jl @@ -205,9 +205,12 @@ for TKey in (Symbol, VarName) ) N = length(samples) D = model.dim + # Follow the sampler's working precision instead of pinning `Float64`, so a + # Float32 (e.g. GPU) run yields a Float32 chain rather than a silently widened one. + FP = typeof(sampler.epsilon) - vals = Matrix{Float64}(undef, N, D) - logp = Vector{Float64}(undef, N) + vals = Matrix{FP}(undef, N, D) + logp = Vector{FP}(undef, N) accepted = Vector{Bool}(undef, N) for i in 1:N @@ -543,7 +546,7 @@ function _parallel_mala_update_progress!( end function _copy_trajectory_rows!( - vals::AbstractMatrix{Float64}, first_row::Int, S::AbstractMatrix, ncols::Int + vals::AbstractMatrix{<:Real}, first_row::Int, S::AbstractMatrix, ncols::Int ) rows = first_row:(first_row + ncols - 1) S_host = Array(view(S, :, 1:ncols)) @@ -563,9 +566,10 @@ function _sample_parallel_mala_chain( progressname="Sampling", ) where {TKey} D = model.dim + FP = typeof(sampler.epsilon) - vals = Matrix{Float64}(undef, N, D) - logp = Vector{Float64}(undef, N) + vals = Matrix{FP}(undef, N, D) + logp = Vector{FP}(undef, N) progress = _parallel_mala_progress(progress, progressname) x0 = _parallel_mala_initial_x(rng, model, sampler, initial_params) @@ -888,9 +892,10 @@ for TKey in (Symbol, VarName) ) N = length(samples) D = model.dim + FP = typeof(sampler.epsilon) - vals = Matrix{Float64}(undef, N, D) - logp = Vector{Float64}(undef, N) + vals = Matrix{FP}(undef, N, D) + logp = Vector{FP}(undef, N) for i in 1:N vals[i, :] .= samples[i].x @@ -1092,11 +1097,12 @@ for TKey in (Symbol, VarName) filtered = discard_warmup ? filter(s -> !s.is_warmup, samples) : samples N = length(filtered) D = model.dim + FP = typeof(sampler.epsilon_init) - vals = Matrix{Float64}(undef, N, D) - logp = Vector{Float64}(undef, N) + vals = Matrix{FP}(undef, N, D) + logp = Vector{FP}(undef, N) accepted = Vector{Bool}(undef, N) - step_size = Vector{Float64}(undef, N) + step_size = Vector{FP}(undef, N) is_warmup = Vector{Bool}(undef, N) for i in 1:N diff --git a/test/test-AbstractMCMC-Interface.jl b/test/test-AbstractMCMC-Interface.jl index c033993..033dc27 100644 --- a/test/test-AbstractMCMC-Interface.jl +++ b/test/test-AbstractMCMC-Interface.jl @@ -152,9 +152,49 @@ gradlogp_iface(x) = -x # logp values should be finite @test all(isfinite, chain[:logp]) - # accepted values should be 0 or 1 + # accepted is stored with its natural `Bool` type (issue #49) acc = chain[:accepted] - @test all(a -> a == 0.0 || a == 1.0, acc) + @test eltype(acc) === Bool + @test all(a -> a == false || a == true, acc) + end + + @testset "chain element type follows sampler precision" begin + # Scalar param names so `chain[:a]` is a scalar-valued column whose eltype is the + # numeric type directly (a vector-valued param would nest inside a `Vector`). + model = DensityModel(logp_iface, gradlogp_iface, 2) + + @testset "Float32 sampler → Float32 chain" begin + sampler = MALASampler(0.1f0) + @test sampler.epsilon isa Float32 + chain = sample( + MersenneTwister(1), + model, + sampler, + 50; + chain_type=SymChain, + progress=false, + param_names=[:a, :b], + ) + @test eltype(chain[:a]) === Float32 + @test eltype(chain[:b]) === Float32 + @test eltype(chain[:logp]) === Float32 + @test eltype(chain[:accepted]) === Bool + end + + @testset "Float64 sampler → Float64 chain" begin + sampler = MALASampler(0.1) + chain = sample( + MersenneTwister(1), + model, + sampler, + 50; + chain_type=SymChain, + progress=false, + param_names=[:a, :b], + ) + @test eltype(chain[:a]) === Float64 + @test eltype(chain[:logp]) === Float64 + end end @testset "sample() with custom param_names" begin diff --git a/test/test-Adaptive-MALA.jl b/test/test-Adaptive-MALA.jl index d4c5a61..5bac2b0 100644 --- a/test/test-Adaptive-MALA.jl +++ b/test/test-Adaptive-MALA.jl @@ -201,13 +201,16 @@ end @test FlexiChains.Extra(:is_warmup) in extras_names @test all(isfinite, chain[:logp]) - @test all(x -> x == 0.0 || x == 1.0, chain[:accepted]) + # accepted / is_warmup keep their natural `Bool` type (issue #49) + @test eltype(chain[:accepted]) === Bool + @test eltype(chain[:is_warmup]) === Bool + @test all(x -> x == false || x == true, chain[:accepted]) @test all(s -> s > 0, chain[:step_size]) # Warmup samples appear at the start - @test chain[:is_warmup][1] == 1.0 - # Post-warmup samples have is_warmup == 0 - @test chain[:is_warmup][end] == 0.0 + @test chain[:is_warmup][1] == true + # Post-warmup samples have is_warmup == false + @test chain[:is_warmup][end] == false end @testset "step_size is constant after warmup" begin From b1d21c6b42d2106d329ca6a5baa6026f83167567 Mon Sep 17 00:00:00 2001 From: Ryan Senne <50930199+rsenne@users.noreply.github.com> Date: Wed, 1 Jul 2026 09:09:00 -0400 Subject: [PATCH 3/4] Review comments --- src/interface.jl | 41 ++++++++++++++--------------------------- 1 file changed, 14 insertions(+), 27 deletions(-) diff --git a/src/interface.jl b/src/interface.jl index e185ead..3713069 100644 --- a/src/interface.jl +++ b/src/interface.jl @@ -621,8 +621,6 @@ function _construct_flexichain( N, D = size(vals) - # Normalize parameter names to `key => dims` pairs (scalars use `()`), so we can - # consume the columns of `vals` uniformly below. if param_names === nothing param_names = model.param_names end @@ -634,40 +632,29 @@ function _construct_flexichain( if n isa Pair to_parameter(n.first) => n.second elseif n isa TKey || n isa Symbol - to_parameter(n) => () + to_parameter(n) else throw(ArgumentError("param_names must be a collection of Pairs, Symbols, or $TKey, got $(typeof(n))")) end end end - # Build a dict of per-key arrays rather than a single stacked array. This lets each - # key keep its natural element type — e.g. `accepted`/`is_warmup` stay `Bool` instead - # of being widened to `Float64` (see issue #49). - data = OrderedDict{FlexiChains.ParameterOrExtra{<:TKey},AbstractArray}() + arr = reshape(vals, N, 1, D) + param_chain = FlexiChains.FlexiChain{TKey}(arr, Tuple(wrapped_param_names)) - offset = 1 - for (key, sz) in wrapped_param_names - ncols = prod(sz) # `prod(()) == 1`, i.e. scalar consumes one column - if sz === () - data[key] = vals[:, offset] - else - cols = offset:(offset + ncols - 1) - data[key] = [collect(reshape(view(vals, i, cols), sz)) for i in 1:N] - end - offset += ncols - end - offset - 1 == D || throw( - ArgumentError( - "param_names consume $(offset - 1) columns, but there are $D parameter columns" - ), + #= + The internals have mixed element types (e.g. `Bool` for `accepted`/`is_warmup`, + `FP` for `logp`), so they cannot share a single stacked array without widening. + Store each in its own column via the dict-of-arrays constructor, which preserves + the natural element type (see issue #49), then merge with the parameters. + =# + extras = OrderedDict{FlexiChains.ParameterOrExtra{<:TKey},AbstractArray}( + FlexiChains.Extra(name) => internals[name] for name in keys(internals) ) + isempty(extras) && return param_chain + extra_chain = FlexiChains.FlexiChain{TKey}(N, 1, extras) - for name in keys(internals) - data[FlexiChains.Extra(name)] = internals[name] - end - - return FlexiChains.FlexiChain{TKey}(N, 1, data) + return merge(param_chain, extra_chain) end function _sample_parallel_mala_blocks( From e0d34a03e80805ab0d7459b3bc3c54c9f3321ca9 Mon Sep 17 00:00:00 2001 From: Ryan Senne <50930199+rsenne@users.noreply.github.com> Date: Wed, 1 Jul 2026 09:18:29 -0400 Subject: [PATCH 4/4] move isempty check --- src/interface.jl | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/src/interface.jl b/src/interface.jl index 3713069..41d1de7 100644 --- a/src/interface.jl +++ b/src/interface.jl @@ -642,6 +642,8 @@ function _construct_flexichain( arr = reshape(vals, N, 1, D) param_chain = FlexiChains.FlexiChain{TKey}(arr, Tuple(wrapped_param_names)) + isempty(internals) && return param_chain + #= The internals have mixed element types (e.g. `Bool` for `accepted`/`is_warmup`, `FP` for `logp`), so they cannot share a single stacked array without widening. @@ -651,7 +653,6 @@ function _construct_flexichain( extras = OrderedDict{FlexiChains.ParameterOrExtra{<:TKey},AbstractArray}( FlexiChains.Extra(name) => internals[name] for name in keys(internals) ) - isempty(extras) && return param_chain extra_chain = FlexiChains.FlexiChain{TKey}(N, 1, extras) return merge(param_chain, extra_chain)