From eed20b0c7f778673a1dbe6aa8c71fec762a11b32 Mon Sep 17 00:00:00 2001
From: Xianda Sun <sunxdt@gmail.com>
Date: Sat, 11 Jul 2026 20:06:15 +0100
Subject: [PATCH 1/5] Allow structured parameters in ParamsWithStats

---
 Project.toml          |  2 +-
 docs/src/callbacks.md |  8 ++++++-
 src/callbacks.jl      | 35 ++++++++++++++++++++++++-----
 test/callbacks.jl     | 51 ++++++++++++++++++++++++++++++++++++++++---
 4 files changed, 85 insertions(+), 11 deletions(-)

diff --git a/Project.toml b/Project.toml
index 7a667fa3..37d9aa42 100644
--- a/Project.toml
+++ b/Project.toml
@@ -3,7 +3,7 @@ uuid = "80f14c24-f653-4e6a-9b94-39d6b0f70001"
 keywords = ["markov chain monte carlo", "probabilistic programming"]
 license = "MIT"
 desc = "A lightweight interface for common MCMC methods."
-version = "5.15.1"
+version = "5.16.0"
 
 [deps]
 BangBang = "198e06fe-97b7-11e9-32a5-e1d131e6ad66"
diff --git a/docs/src/callbacks.md b/docs/src/callbacks.md
index 0156f238..217e3758 100644
--- a/docs/src/callbacks.md
+++ b/docs/src/callbacks.md
@@ -250,11 +250,17 @@ function AbstractMCMC.getstats(state::MyState)
 end
 ```
 
-The `ParamsWithStats` constructors normalize all inputs to `NamedTuple`:
+The default `ParamsWithStats` extraction constructors normalize common inputs to
+`NamedTuple`:
 - `Vector{<:Real}` gets default `θ[i]` names
 - `Vector{Pair}` is converted to `NamedTuple` with the provided names
 - `NamedTuple` is used directly
 
+Packages with a structured parameter representation can also construct
+`ParamsWithStats(params, stats)` directly. The parameter container must implement `pairs`
+and `isempty`; its keys do not need to be symbols. The two-argument constructor uses an
+empty `NamedTuple` for `extras`.
+
 !!! note "stats vs extras"
     Use `stats` for values that change once per MCMC iteration (e.g., log probability, acceptance rate).
     Use `extras` for values that are constant across iterations (e.g., preconditioning matrix, number of particles)
diff --git a/src/callbacks.jl b/src/callbacks.jl
index 9c43b17a..e347e36e 100644
--- a/src/callbacks.jl
+++ b/src/callbacks.jl
@@ -118,13 +118,12 @@ end
 """
     ParamsWithStats{P,S,E}
 
-A container for MCMC parameters, statistics, and extras.
-
-All fields are stored as `NamedTuple`s to ensure a tight, well-defined interface.
-Use `Base.pairs(pws)` to iterate over `(name, value)` pairs.
+A container for MCMC parameters, statistics, and extras. The parameter container can be any
+type implementing `pairs` and `isempty`; statistics and extras are stored as `NamedTuple`s.
+Use `Base.pairs(pws)` to iterate over all `(name, value)` pairs.
 
 # Fields
-- `params::P`: Parameter values as a NamedTuple
+- `params::P`: Parameter values in a container implementing `pairs` and `isempty`
 - `stats::S`: Statistics as a NamedTuple (e.g., `(lp=...,)`)
 - `extras::E`: Extra diagnostics as a NamedTuple
 
@@ -139,12 +138,21 @@ end
 pws2 = ParamsWithStats(pws; params=true, stats=false)
 ```
 """
-struct ParamsWithStats{P<:NamedTuple,S<:NamedTuple,E<:NamedTuple}
+struct ParamsWithStats{P,S<:NamedTuple,E<:NamedTuple}
     params::P
     stats::S
     extras::E
 end
 
+"""
+    ParamsWithStats(params, stats::NamedTuple)
+
+Construct a `ParamsWithStats` with no extra diagnostics.
+"""
+function ParamsWithStats(params, stats::NamedTuple)
+    return ParamsWithStats(params, stats, NamedTuple())
+end
+
 # Constructor from Vector{<:Real} - adds default θ[i] names
 function ParamsWithStats(
     v::AbstractVector{<:Real}, stats::S, extras::E
@@ -237,6 +245,21 @@ function Base.isempty(pws::ParamsWithStats)
     return (isempty(pws.params) && isempty(pws.stats) && isempty(pws.extras))
 end
 
+function Base.:(==)(pws1::ParamsWithStats, pws2::ParamsWithStats)
+    return (pws1.params == pws2.params) & (pws1.stats == pws2.stats) &
+           (pws1.extras == pws2.extras)
+end
+
+function Base.isequal(pws1::ParamsWithStats, pws2::ParamsWithStats)
+    return isequal(pws1.params, pws2.params) &&
+           isequal(pws1.stats, pws2.stats) &&
+           isequal(pws1.extras, pws2.extras)
+end
+
+function Base.hash(pws::ParamsWithStats, h::UInt)
+    return hash(pws.extras, hash(pws.stats, hash(pws.params, hash(:ParamsWithStats, h))))
+end
+
 #################################
 ### Unified mcmc_callback API ###
 #################################
diff --git a/test/callbacks.jl b/test/callbacks.jl
index 24903b55..749e2d1e 100644
--- a/test/callbacks.jl
+++ b/test/callbacks.jl
@@ -166,6 +166,15 @@ end
 ### ParamsWithStats   ###
 #########################
 
+struct CustomParams{T}
+    data::T
+end
+Base.pairs(params::CustomParams) = pairs(params.data)
+Base.isempty(params::CustomParams) = isempty(params.data)
+Base.:(==)(a::CustomParams, b::CustomParams) = a.data == b.data
+Base.isequal(a::CustomParams, b::CustomParams) = isequal(a.data, b.data)
+Base.hash(params::CustomParams, h::UInt) = hash(params.data, h)
+
 @testset "ParamsWithStats" begin
     @testset "Constructor from NamedTuple" begin
         pws = AbstractMCMC.ParamsWithStats((a=1.0, b=2.0), (lp=-10.0,), NamedTuple())
@@ -178,6 +187,11 @@ end
     @testset "Constructor from Vector{Real} - default names" begin
         pws = AbstractMCMC.ParamsWithStats([1.0, 2.0, 3.0], NamedTuple(), NamedTuple())
         @test pws.params == (var"θ[1]"=1.0, var"θ[2]"=2.0, var"θ[3]"=3.0)
+
+        pws_without_extras = AbstractMCMC.ParamsWithStats([1.0, 2.0], (lp=-1.0,))
+        @test pws_without_extras.params == (var"θ[1]"=1.0, var"θ[2]"=2.0)
+        @test pws_without_extras.stats == (lp=-1.0,)
+        @test isempty(pws_without_extras.extras)
     end
 
     @testset "Constructor from Vector{Pair} - named" begin
@@ -187,6 +201,19 @@ end
         @test pws.params == (μ=1.0, σ=2.0)
     end
 
+    @testset "Constructor from custom parameter container" begin
+        params = CustomParams((x=[1.0, 2.0], y=3.0))
+        pws = AbstractMCMC.ParamsWithStats(params, (lp=-10.0,))
+        @test pws.params === params
+        @test pws.stats == (lp=-10.0,)
+        @test pws.extras == NamedTuple()
+        @test collect(pairs(pws)) == [:x => [1.0, 2.0], :y => 3.0, :lp => -10.0]
+        @test !isempty(pws)
+
+        empty_pws = AbstractMCMC.ParamsWithStats(CustomParams(NamedTuple()), NamedTuple())
+        @test isempty(empty_pws)
+    end
+
     @testset "Constructor from state" begin
         state = 5
         pws = AbstractMCMC.ParamsWithStats(
@@ -210,6 +237,11 @@ end
         pws_stats = AbstractMCMC.ParamsWithStats(pws; params=false, stats=true)
         @test pws_stats.params == NamedTuple()
         @test pws_stats.stats == (lp=-10.0,)
+
+        custom_params = CustomParams((x=1.0,))
+        custom_pws = AbstractMCMC.ParamsWithStats(custom_params, NamedTuple())
+        @test AbstractMCMC.ParamsWithStats(custom_pws; stats=false).params === custom_params
+        @test AbstractMCMC.ParamsWithStats(custom_pws; params=false).params == NamedTuple()
     end
 
     @testset "Base.pairs iteration" begin
@@ -230,11 +262,24 @@ end
     end
 
     @testset "Illegal states are unrepresentable" begin
-        # Should not be able to construct with arbitrary types
+        # Statistics and extras must be NamedTuples.
         @test_throws MethodError AbstractMCMC.ParamsWithStats(1, 2, 3)
-        @test_throws MethodError AbstractMCMC.ParamsWithStats(
-            "bad", NamedTuple(), NamedTuple()
+        @test_throws MethodError AbstractMCMC.ParamsWithStats("params", 2, NamedTuple())
+        @test_throws MethodError AbstractMCMC.ParamsWithStats("params", NamedTuple(), 3)
+    end
+
+    @testset "Equality" begin
+        pws1 = AbstractMCMC.ParamsWithStats(
+            CustomParams((x=[1.0, NaN],)), (lp=-10.0,), (step_size=0.1,)
+        )
+        pws2 = AbstractMCMC.ParamsWithStats(
+            CustomParams((x=[1.0, NaN],)), (lp=-10.0,), (step_size=0.1,)
         )
+        @test isequal(pws1, pws2)
+        @test hash(pws1) == hash(pws2)
+        @test !(pws1 == pws2)
+        @test pws1 !=
+            AbstractMCMC.ParamsWithStats(pws1.params, pws1.stats, (step_size=0.2,))
     end
 end
 

From a54693d76ac3b4edb2edc9c7116ce67573dae602 Mon Sep 17 00:00:00 2001
From: Xianda Sun <sunxdt@gmail.com>
Date: Mon, 13 Jul 2026 09:09:26 +0100
Subject: [PATCH 2/5] Document the full params-container contract and mixed-key
 pairs iteration

---
 docs/src/callbacks.md | 14 ++++++++++++--
 src/callbacks.jl      | 17 ++++++++++++++---
 2 files changed, 26 insertions(+), 5 deletions(-)

diff --git a/docs/src/callbacks.md b/docs/src/callbacks.md
index 217e3758..344d6dc6 100644
--- a/docs/src/callbacks.md
+++ b/docs/src/callbacks.md
@@ -258,8 +258,18 @@ The default `ParamsWithStats` extraction constructors normalize common inputs to
 
 Packages with a structured parameter representation can also construct
 `ParamsWithStats(params, stats)` directly. The parameter container must implement `pairs`
-and `isempty`; its keys do not need to be symbols. The two-argument constructor uses an
-empty `NamedTuple` for `extras`.
+and `isempty`; for `==`, `isequal`, and `hash` of the wrapper to be meaningful it should
+implement those as well, and its keys should have a meaningful `string` form so name-based
+filtering and logging work. The two-argument constructor uses an empty `NamedTuple` for
+`extras`. Note the normalization above applies to `AbstractVector` subtypes: a vector-like
+container is converted to a `Symbol`-keyed `NamedTuple` rather than stored as given.
+
+!!! warning "Mixed key types"
+    Because the parameter container's keys need not be `Symbol`s, `Base.pairs(pws)` may
+    yield pairs with mixed key types (e.g., `VarName` keys from the parameters and
+    `Symbol` keys from the statistics). Consumers should iterate the pairs generically
+    (e.g., stringify keys) and must not assume `Symbol` keys, a concrete element type, or
+    that the pairs can be collected into a `NamedTuple`.
 
 !!! note "stats vs extras"
     Use `stats` for values that change once per MCMC iteration (e.g., log probability, acceptance rate).
diff --git a/src/callbacks.jl b/src/callbacks.jl
index e347e36e..e36c3f88 100644
--- a/src/callbacks.jl
+++ b/src/callbacks.jl
@@ -118,9 +118,20 @@ end
 """
     ParamsWithStats{P,S,E}
 
-A container for MCMC parameters, statistics, and extras. The parameter container can be any
-type implementing `pairs` and `isempty`; statistics and extras are stored as `NamedTuple`s.
-Use `Base.pairs(pws)` to iterate over all `(name, value)` pairs.
+A container for MCMC parameters, statistics, and extras. The parameter container can be a
+structured type; statistics and extras are stored as `NamedTuple`s. Use `Base.pairs(pws)`
+to iterate over all `(name, value)` pairs.
+
+The parameter container must implement `pairs` and `isempty`. For `==`, `isequal`, and
+`hash` of `ParamsWithStats` to be meaningful it must also implement those (with `==`
+returning `Bool` or `missing`), and its keys should have a meaningful `string` form so
+that name-based filtering and logging callbacks work. Keys are not required to be
+`Symbol`s, so `pairs(pws)` may yield pairs with mixed key types; consumers should not
+assume `Symbol` keys or a concrete element type.
+
+Note that `AbstractVector{<:Real}` and `AbstractVector{<:Pair}` parameter inputs are not
+stored as given: the extraction constructors normalize them to `Symbol`-keyed
+`NamedTuple`s (see the constructor docs below).
 
 # Fields
 - `params::P`: Parameter values in a container implementing `pairs` and `isempty`

From 4f13c697b733d26450193c09c24f38cffa95438a Mon Sep 17 00:00:00 2001
From: Hong Ge <hg344@cam.ac.uk>
Date: Tue, 21 Jul 2026 15:35:01 +0100
Subject: [PATCH 3/5] Isolate extras-sensitivity in ParamsWithStats equality
 test
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

The final assertion in the Equality testset used NaN-containing params, so
the params term already forced `==` to false regardless of extras — it never
actually tested that a differing field makes two wrappers unequal. Replace it
with a NaN-free comparison plus a matching equal case, and normalise `==`
formatting to blue style.

Co-Authored-By: Claude Code <noreply@anthropic.com>
---
 src/callbacks.jl  |  3 ++-
 test/callbacks.jl | 14 ++++++++++++--
 2 files changed, 14 insertions(+), 3 deletions(-)

diff --git a/src/callbacks.jl b/src/callbacks.jl
index e36c3f88..84dbc11a 100644
--- a/src/callbacks.jl
+++ b/src/callbacks.jl
@@ -257,7 +257,8 @@ function Base.isempty(pws::ParamsWithStats)
 end
 
 function Base.:(==)(pws1::ParamsWithStats, pws2::ParamsWithStats)
-    return (pws1.params == pws2.params) & (pws1.stats == pws2.stats) &
+    return (pws1.params == pws2.params) &
+           (pws1.stats == pws2.stats) &
            (pws1.extras == pws2.extras)
 end
 
diff --git a/test/callbacks.jl b/test/callbacks.jl
index 749e2d1e..b2312e2b 100644
--- a/test/callbacks.jl
+++ b/test/callbacks.jl
@@ -275,11 +275,21 @@ Base.hash(params::CustomParams, h::UInt) = hash(params.data, h)
         pws2 = AbstractMCMC.ParamsWithStats(
             CustomParams((x=[1.0, NaN],)), (lp=-10.0,), (step_size=0.1,)
         )
+        # NaN params: isequal/hash match, but `==` is false (NaN != NaN).
         @test isequal(pws1, pws2)
         @test hash(pws1) == hash(pws2)
         @test !(pws1 == pws2)
-        @test pws1 !=
-            AbstractMCMC.ParamsWithStats(pws1.params, pws1.stats, (step_size=0.2,))
+
+        # A differing field alone makes `==` false (NaN-free, so params don't confound).
+        base = AbstractMCMC.ParamsWithStats(
+            CustomParams((x=1.0,)), (lp=-10.0,), (step_size=0.1,)
+        )
+        @test base == AbstractMCMC.ParamsWithStats(
+            CustomParams((x=1.0,)), (lp=-10.0,), (step_size=0.1,)
+        )
+        @test base != AbstractMCMC.ParamsWithStats(
+            CustomParams((x=1.0,)), (lp=-10.0,), (step_size=0.2,)
+        )
     end
 end
 

From 4cf5d9d24d2dc90d1a957a2b8efffa01ab54fd07 Mon Sep 17 00:00:00 2001
From: Xianda Sun <sunxdt@gmail.com>
Date: Wed, 22 Jul 2026 09:29:39 +0100
Subject: [PATCH 4/5] Drop ==/isequal/hash from ParamsWithStats; contract is
 pairs/isempty only

The wrapper keeps the default struct equality it had before this branch,
so the parameter container only needs to implement pairs and isempty.
Tests now use a VarName-like stand-in (StructuredKey/StructuredParams)
with non-Symbol keys and cover mixed-key iteration, NameFilter on
stringified keys, and container preservation through the extraction
constructor.
---
 docs/src/callbacks.md | 11 +++---
 src/callbacks.jl      | 31 +++++----------
 test/callbacks.jl     | 91 +++++++++++++++++++++++++------------------
 3 files changed, 69 insertions(+), 64 deletions(-)

diff --git a/docs/src/callbacks.md b/docs/src/callbacks.md
index 344d6dc6..3ca93714 100644
--- a/docs/src/callbacks.md
+++ b/docs/src/callbacks.md
@@ -258,11 +258,12 @@ The default `ParamsWithStats` extraction constructors normalize common inputs to
 
 Packages with a structured parameter representation can also construct
 `ParamsWithStats(params, stats)` directly. The parameter container must implement `pairs`
-and `isempty`; for `==`, `isequal`, and `hash` of the wrapper to be meaningful it should
-implement those as well, and its keys should have a meaningful `string` form so name-based
-filtering and logging work. The two-argument constructor uses an empty `NamedTuple` for
-`extras`. Note the normalization above applies to `AbstractVector` subtypes: a vector-like
-container is converted to a `Symbol`-keyed `NamedTuple` rather than stored as given.
+and `isempty`; this is not validated at construction (some Base types already define those
+methods and will not fail later). Keys should have a meaningful `string` form so
+name-based filtering and logging work. The two-argument constructor uses an empty
+`NamedTuple` for `extras`. Note the normalization above applies to `AbstractVector`
+subtypes: a vector-like container is converted to a `Symbol`-keyed `NamedTuple` rather
+than stored as given.
 
 !!! warning "Mixed key types"
     Because the parameter container's keys need not be `Symbol`s, `Base.pairs(pws)` may
diff --git a/src/callbacks.jl b/src/callbacks.jl
index 84dbc11a..836c3c5e 100644
--- a/src/callbacks.jl
+++ b/src/callbacks.jl
@@ -122,12 +122,11 @@ A container for MCMC parameters, statistics, and extras. The parameter container
 structured type; statistics and extras are stored as `NamedTuple`s. Use `Base.pairs(pws)`
 to iterate over all `(name, value)` pairs.
 
-The parameter container must implement `pairs` and `isempty`. For `==`, `isequal`, and
-`hash` of `ParamsWithStats` to be meaningful it must also implement those (with `==`
-returning `Bool` or `missing`), and its keys should have a meaningful `string` form so
-that name-based filtering and logging callbacks work. Keys are not required to be
-`Symbol`s, so `pairs(pws)` may yield pairs with mixed key types; consumers should not
-assume `Symbol` keys or a concrete element type.
+The parameter container must implement `pairs` and `isempty`. This is not validated at
+construction; some Base types already define those methods and will not fail later. Keys
+should have a meaningful `string` form so that name-based filtering and logging callbacks
+work. Keys are not required to be `Symbol`s, so `pairs(pws)` may yield pairs with mixed
+key types; consumers should not assume `Symbol` keys or a concrete element type.
 
 Note that `AbstractVector{<:Real}` and `AbstractVector{<:Pair}` parameter inputs are not
 stored as given: the extraction constructors normalize them to `Symbol`-keyed
@@ -241,6 +240,10 @@ end
 
 Return an iterator of `(name, value)` pairs for all selected data in `pws`.
 
+Parameter keys need not be `Symbol`s, so the iterator may yield mixed key types (e.g.
+structured parameter keys together with `Symbol` keys from `stats`/`extras`). Consumers
+should iterate generically and must not assume `Symbol` keys or a concrete element type.
+
 This is the canonical way to iterate over a `ParamsWithStats`:
 ```julia
 for (name, value) in Base.pairs(pws)
@@ -256,22 +259,6 @@ function Base.isempty(pws::ParamsWithStats)
     return (isempty(pws.params) && isempty(pws.stats) && isempty(pws.extras))
 end
 
-function Base.:(==)(pws1::ParamsWithStats, pws2::ParamsWithStats)
-    return (pws1.params == pws2.params) &
-           (pws1.stats == pws2.stats) &
-           (pws1.extras == pws2.extras)
-end
-
-function Base.isequal(pws1::ParamsWithStats, pws2::ParamsWithStats)
-    return isequal(pws1.params, pws2.params) &&
-           isequal(pws1.stats, pws2.stats) &&
-           isequal(pws1.extras, pws2.extras)
-end
-
-function Base.hash(pws::ParamsWithStats, h::UInt)
-    return hash(pws.extras, hash(pws.stats, hash(pws.params, hash(:ParamsWithStats, h))))
-end
-
 #################################
 ### Unified mcmc_callback API ###
 #################################
diff --git a/test/callbacks.jl b/test/callbacks.jl
index b2312e2b..059f279f 100644
--- a/test/callbacks.jl
+++ b/test/callbacks.jl
@@ -166,14 +166,19 @@ end
 ### ParamsWithStats   ###
 #########################
 
-struct CustomParams{T}
-    data::T
+# StructuredKey: non-Symbol key with string form (VarName stand-in). StructuredParams
+# wraps a Vector of Pairs so AbstractVector{<:Pair} normalization does not kick in.
+struct StructuredKey
+    name::String
 end
-Base.pairs(params::CustomParams) = pairs(params.data)
-Base.isempty(params::CustomParams) = isempty(params.data)
-Base.:(==)(a::CustomParams, b::CustomParams) = a.data == b.data
-Base.isequal(a::CustomParams, b::CustomParams) = isequal(a.data, b.data)
-Base.hash(params::CustomParams, h::UInt) = hash(params.data, h)
+Base.string(k::StructuredKey) = k.name
+Base.:(==)(a::StructuredKey, b::StructuredKey) = a.name == b.name
+
+struct StructuredParams{P<:Pair}
+    data::Vector{P}
+end
+Base.pairs(params::StructuredParams) = params.data
+Base.isempty(params::StructuredParams) = isempty(params.data)
 
 @testset "ParamsWithStats" begin
     @testset "Constructor from NamedTuple" begin
@@ -202,18 +207,35 @@ Base.hash(params::CustomParams, h::UInt) = hash(params.data, h)
     end
 
     @testset "Constructor from custom parameter container" begin
-        params = CustomParams((x=[1.0, 2.0], y=3.0))
+        params = StructuredParams([StructuredKey("x") => 1.0, StructuredKey("y") => 3.0])
         pws = AbstractMCMC.ParamsWithStats(params, (lp=-10.0,))
         @test pws.params === params
         @test pws.stats == (lp=-10.0,)
         @test pws.extras == NamedTuple()
-        @test collect(pairs(pws)) == [:x => [1.0, 2.0], :y => 3.0, :lp => -10.0]
+        @test collect(pairs(pws)) ==
+            [StructuredKey("x") => 1.0, StructuredKey("y") => 3.0, :lp => -10.0]
         @test !isempty(pws)
 
-        empty_pws = AbstractMCMC.ParamsWithStats(CustomParams(NamedTuple()), NamedTuple())
+        empty_pws = AbstractMCMC.ParamsWithStats(
+            StructuredParams(Pair{StructuredKey,Float64}[]), NamedTuple()
+        )
         @test isempty(empty_pws)
     end
 
+    @testset "Mixed key types from params and stats" begin
+        params = StructuredParams([StructuredKey("μ") => 1.0])
+        pws = AbstractMCMC.ParamsWithStats(params, (lp=-10.0,), (step_size=0.1,))
+        pairs_list = collect(Base.pairs(pws))
+        @test pairs_list == [StructuredKey("μ") => 1.0, :lp => -10.0, :step_size => 0.1]
+        @test pairs_list[1][1] isa StructuredKey
+        @test pairs_list[2][1] isa Symbol
+
+        # NameFilter stringifies keys, so non-Symbol parameter keys still filter.
+        f = AbstractMCMC.NameFilter(; include=["μ", "lp"])
+        filtered = collect(Iterators.filter(((k, _),) -> f(k), pairs_list))
+        @test filtered == [StructuredKey("μ") => 1.0, :lp => -10.0]
+    end
+
     @testset "Constructor from state" begin
         state = 5
         pws = AbstractMCMC.ParamsWithStats(
@@ -225,6 +247,22 @@ Base.hash(params::CustomParams, h::UInt) = hash(params.data, h)
         @test pws.extras == NamedTuple()
     end
 
+    @testset "Extraction constructor preserves custom parameter container" begin
+        struct StructuredState
+            params::StructuredParams
+        end
+        AbstractMCMC.getparams(s::StructuredState) = s.params
+        AbstractMCMC.getstats(s::StructuredState) = (lp=-3.0,)
+
+        params = StructuredParams([StructuredKey("α") => 0.5])
+        state = StructuredState(params)
+        pws = AbstractMCMC.ParamsWithStats(
+            MyModel(), MySampler(), nothing, state; params=true, stats=true
+        )
+        @test pws.params === params
+        @test collect(Base.pairs(pws)) == [StructuredKey("α") => 0.5, :lp => -3.0]
+    end
+
     @testset "Copy constructor with selection" begin
         pws = AbstractMCMC.ParamsWithStats((a=1.0,), (lp=-10.0,), NamedTuple())
 
@@ -238,7 +276,7 @@ Base.hash(params::CustomParams, h::UInt) = hash(params.data, h)
         @test pws_stats.params == NamedTuple()
         @test pws_stats.stats == (lp=-10.0,)
 
-        custom_params = CustomParams((x=1.0,))
+        custom_params = StructuredParams([StructuredKey("x") => 1.0])
         custom_pws = AbstractMCMC.ParamsWithStats(custom_params, NamedTuple())
         @test AbstractMCMC.ParamsWithStats(custom_pws; stats=false).params === custom_params
         @test AbstractMCMC.ParamsWithStats(custom_pws; params=false).params == NamedTuple()
@@ -261,35 +299,14 @@ Base.hash(params::CustomParams, h::UInt) = hash(params.data, h)
         @test isempty(pws_empty)
     end
 
-    @testset "Illegal states are unrepresentable" begin
-        # Statistics and extras must be NamedTuples.
+    @testset "Stats and extras are constrained" begin
+        # Statistics and extras must be NamedTuples. Params are not validated at
+        # construction (e.g. `1` still constructs; Base defines pairs/isempty for it).
         @test_throws MethodError AbstractMCMC.ParamsWithStats(1, 2, 3)
         @test_throws MethodError AbstractMCMC.ParamsWithStats("params", 2, NamedTuple())
         @test_throws MethodError AbstractMCMC.ParamsWithStats("params", NamedTuple(), 3)
-    end
-
-    @testset "Equality" begin
-        pws1 = AbstractMCMC.ParamsWithStats(
-            CustomParams((x=[1.0, NaN],)), (lp=-10.0,), (step_size=0.1,)
-        )
-        pws2 = AbstractMCMC.ParamsWithStats(
-            CustomParams((x=[1.0, NaN],)), (lp=-10.0,), (step_size=0.1,)
-        )
-        # NaN params: isequal/hash match, but `==` is false (NaN != NaN).
-        @test isequal(pws1, pws2)
-        @test hash(pws1) == hash(pws2)
-        @test !(pws1 == pws2)
-
-        # A differing field alone makes `==` false (NaN-free, so params don't confound).
-        base = AbstractMCMC.ParamsWithStats(
-            CustomParams((x=1.0,)), (lp=-10.0,), (step_size=0.1,)
-        )
-        @test base == AbstractMCMC.ParamsWithStats(
-            CustomParams((x=1.0,)), (lp=-10.0,), (step_size=0.1,)
-        )
-        @test base != AbstractMCMC.ParamsWithStats(
-            CustomParams((x=1.0,)), (lp=-10.0,), (step_size=0.2,)
-        )
+        @test AbstractMCMC.ParamsWithStats(1, NamedTuple(), NamedTuple()) isa
+            AbstractMCMC.ParamsWithStats
     end
 end
 

From 31e296ea91d3c548e6ef36edf96df2b5a5652c78 Mon Sep 17 00:00:00 2001
From: Xianda Sun <sunxdt@gmail.com>
Date: Wed, 22 Jul 2026 09:44:02 +0100
Subject: [PATCH 5/5] Consolidate structured-container tests

Merge the overlapping structured-params testsets, drop the StructuredState
wrapper by letting the container act as its own state, and simplify
StructuredParams to a concrete field. Same code paths covered with one
setup per testset.
---
 test/callbacks.jl | 57 +++++++++++++++++++++--------------------------
 1 file changed, 26 insertions(+), 31 deletions(-)

diff --git a/test/callbacks.jl b/test/callbacks.jl
index 059f279f..8f1f2821 100644
--- a/test/callbacks.jl
+++ b/test/callbacks.jl
@@ -174,8 +174,8 @@ end
 Base.string(k::StructuredKey) = k.name
 Base.:(==)(a::StructuredKey, b::StructuredKey) = a.name == b.name
 
-struct StructuredParams{P<:Pair}
-    data::Vector{P}
+struct StructuredParams
+    data::Vector{Pair{StructuredKey,Float64}}
 end
 Base.pairs(params::StructuredParams) = params.data
 Base.isempty(params::StructuredParams) = isempty(params.data)
@@ -206,33 +206,32 @@ Base.isempty(params::StructuredParams) = isempty(params.data)
         @test pws.params == (μ=1.0, σ=2.0)
     end
 
-    @testset "Constructor from custom parameter container" begin
-        params = StructuredParams([StructuredKey("x") => 1.0, StructuredKey("y") => 3.0])
+    @testset "Structured parameter container" begin
+        params = StructuredParams([StructuredKey("μ") => 1.0, StructuredKey("y") => 3.0])
+
+        # The two-argument constructor stores the container as-is and defaults extras.
         pws = AbstractMCMC.ParamsWithStats(params, (lp=-10.0,))
         @test pws.params === params
         @test pws.stats == (lp=-10.0,)
         @test pws.extras == NamedTuple()
-        @test collect(pairs(pws)) ==
-            [StructuredKey("x") => 1.0, StructuredKey("y") => 3.0, :lp => -10.0]
-        @test !isempty(pws)
-
-        empty_pws = AbstractMCMC.ParamsWithStats(
-            StructuredParams(Pair{StructuredKey,Float64}[]), NamedTuple()
-        )
-        @test isempty(empty_pws)
-    end
 
-    @testset "Mixed key types from params and stats" begin
-        params = StructuredParams([StructuredKey("μ") => 1.0])
+        # pairs flattens params/stats/extras and yields mixed key types.
         pws = AbstractMCMC.ParamsWithStats(params, (lp=-10.0,), (step_size=0.1,))
         pairs_list = collect(Base.pairs(pws))
-        @test pairs_list == [StructuredKey("μ") => 1.0, :lp => -10.0, :step_size => 0.1]
+        @test pairs_list == [
+            StructuredKey("μ") => 1.0,
+            StructuredKey("y") => 3.0,
+            :lp => -10.0,
+            :step_size => 0.1,
+        ]
         @test pairs_list[1][1] isa StructuredKey
-        @test pairs_list[2][1] isa Symbol
+        @test pairs_list[3][1] isa Symbol
+        @test !isempty(pws)
+        @test isempty(AbstractMCMC.ParamsWithStats(StructuredParams([]), NamedTuple()))
 
         # NameFilter stringifies keys, so non-Symbol parameter keys still filter.
         f = AbstractMCMC.NameFilter(; include=["μ", "lp"])
-        filtered = collect(Iterators.filter(((k, _),) -> f(k), pairs_list))
+        filtered = filter(p -> f(first(p)), pairs_list)
         @test filtered == [StructuredKey("μ") => 1.0, :lp => -10.0]
     end
 
@@ -247,20 +246,21 @@ Base.isempty(params::StructuredParams) = isempty(params.data)
         @test pws.extras == NamedTuple()
     end
 
-    @testset "Extraction constructor preserves custom parameter container" begin
-        struct StructuredState
-            params::StructuredParams
-        end
-        AbstractMCMC.getparams(s::StructuredState) = s.params
-        AbstractMCMC.getstats(s::StructuredState) = (lp=-3.0,)
+    @testset "Structured container through extraction and re-selection" begin
+        # Let the container act as its own state so getparams/getstats drive extraction.
+        AbstractMCMC.getparams(params::StructuredParams) = params
+        AbstractMCMC.getstats(::StructuredParams) = (lp=-3.0,)
 
         params = StructuredParams([StructuredKey("α") => 0.5])
-        state = StructuredState(params)
         pws = AbstractMCMC.ParamsWithStats(
-            MyModel(), MySampler(), nothing, state; params=true, stats=true
+            MyModel(), MySampler(), nothing, params; params=true, stats=true
         )
         @test pws.params === params
         @test collect(Base.pairs(pws)) == [StructuredKey("α") => 0.5, :lp => -3.0]
+
+        # Re-selection preserves the container; deselecting substitutes an empty NamedTuple.
+        @test AbstractMCMC.ParamsWithStats(pws; stats=false).params === params
+        @test AbstractMCMC.ParamsWithStats(pws; params=false).params == NamedTuple()
     end
 
     @testset "Copy constructor with selection" begin
@@ -275,11 +275,6 @@ Base.isempty(params::StructuredParams) = isempty(params.data)
         pws_stats = AbstractMCMC.ParamsWithStats(pws; params=false, stats=true)
         @test pws_stats.params == NamedTuple()
         @test pws_stats.stats == (lp=-10.0,)
-
-        custom_params = StructuredParams([StructuredKey("x") => 1.0])
-        custom_pws = AbstractMCMC.ParamsWithStats(custom_params, NamedTuple())
-        @test AbstractMCMC.ParamsWithStats(custom_pws; stats=false).params === custom_params
-        @test AbstractMCMC.ParamsWithStats(custom_pws; params=false).params == NamedTuple()
     end
 
     @testset "Base.pairs iteration" begin