From b81d463c0ab6748a60e17847bb81145e6a627fda Mon Sep 17 00:00:00 2001
From: Ryan Senne <50930199+rsenne@users.noreply.github.com>
Date: Sat, 24 Oct 2026 13:54:39 -0400
Subject: [PATCH] Bump FlexiChains compat to 0.7 and widen CUDA to 6

---
 CHANGELOG.md                        |  7 +++++++
 Project.toml                        |  4 ++--
 docs/src/10-getting-started.md      |  2 +-
 docs/src/index.md                   |  4 ++--
 src/interface.jl                    | 29 +++++++++++++++++++++++++++++
 test/test-AbstractMCMC-Interface.jl | 15 ++++++++-------
 test/test-Adaptive-MALA.jl          |  5 +++--
 test/test-DEER-Interface.jl         | 14 +++++++++++---
 test/test-DEER-Logistic.jl          | 10 +++++++++-
 test/test-GPU-AD-HVP.jl             |  3 +++
 test/test-GPU-DEER.jl               |  1 +
 test/test-GPU-Performance.jl        |  9 ++++++++-
 test/test-LogDensityProblems.jl     |  4 ++--
 13 files changed, 86 insertions(+), 21 deletions(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index a98d062..c41a873 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -35,6 +35,13 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ### Changed
 
+- FlexiChains compat is now `0.7`. FlexiChains 0.7 takes its `VarName` from
+  VarNames.jl, so `VNChain` is a `FlexiChain{VarNames.VarName}`. The
+  `@varname` macro is no longer exported there; use `@vn`.
+- `sample` on a `DensityModel` now returns a `VNChain` by default. Pass
+  `chain_type=Any` for the raw transition vector, or `chain_type=SymChain`
+  for `Symbol` keys.
+- CUDA compat widened to `5.11, 6`.
 - CUDA is now a weak dependency behind the `CUDAExt` extension; `using CUDA`
   loads it (#59). The `CuArray` checks in the random fills are replaced by the
   `needs_host_staging` trait.
diff --git a/Project.toml b/Project.toml
index ac2df19..f6eb613 100644
--- a/Project.toml
+++ b/Project.toml
@@ -31,11 +31,11 @@ ReactantExt = ["Reactant", "Enzyme"]
 [compat]
 ADTypes = "1.21.0"
 AbstractMCMC = "5.10.0"
-CUDA = "5.11.0"
+CUDA = "5.11, 6"
 CUDA_Runtime_jll = "0.21"
 DifferentiationInterface = "0.7.13"
 Enzyme = "0.13.146"
-FlexiChains = "0.6.6"
+FlexiChains = "0.7"
 ForwardDiff = "1"
 LinearAlgebra = "1"
 LogDensityProblems = "2"
diff --git a/docs/src/10-getting-started.md b/docs/src/10-getting-started.md
index 7475b75..a250043 100644
--- a/docs/src/10-getting-started.md
+++ b/docs/src/10-getting-started.md
@@ -171,7 +171,7 @@ If you do not specify `param_names`, the chain will store a single vector-valued
 You can specify parameter names as a collection of either:
 
 - `Symbol`s (e.g. `[:x1, :x2]`);
-- `VarName`s (e.g. `[@varname(x[1]), @varname(x[2])]`) if `chain_type=VNChain`; or
+- `VarName`s (e.g. `[@vn(x[1]), @vn(x[2])]`) if `chain_type=VNChain`; or
 - A tuple of the above, *plus* a size. In this case, a total of `prod(size)` entries will be allocated to the named parameter, and the results in the chain will be reshaped to that size.
 
 The above can be mixed and matched as desired, as long as the total number of parameters matches the dimension of the model.
diff --git a/docs/src/index.md b/docs/src/index.md
index 0bd8784..e304f3d 100644
--- a/docs/src/index.md
+++ b/docs/src/index.md
@@ -80,8 +80,8 @@ chain = sample(model, sampler, 500;
 
 Each call to `sample` draws 500 samples by solving DEER trajectories of length `T=64` in parallel, re-solving from the last state when each trajectory is exhausted.
 
-Specifying `chain_type=VNChain` returns a `FlexiChain{VarName}`, which has a parameter type of `VarName`.
-This is intended for maximum ease of use; however, if you prefer parameter type of `Symbol` you can use `SymChain` instead.
+`sample` returns a `VNChain`, a `FlexiChain` keyed by `VarName`, so the `chain_type=VNChain` above is the default and can be left out.
+`SymChain` (keyed by `Symbol`) and `Any` (the raw transition vector) are still accepted but not recommended.
 See [the FlexiChains.jl docs](https://pysm.dev/FlexiChains.jl/) for more information about how to analyze and visualize chains.
 
 ### Sequential MALA baseline
diff --git a/src/interface.jl b/src/interface.jl
index 8935f84..5aa0523 100644
--- a/src/interface.jl
+++ b/src/interface.jl
@@ -1515,3 +1515,32 @@ for TKey in (Symbol, VarName)
         return _construct_flexichain($TKey, vals, internals, param_names, model)
     end
 end
+
+#= `sample` on a `DensityModel` returns a `VNChain` unless `chain_type` says
+otherwise. AbstractMCMC's own default is `Any`, which hands back the raw
+transition vector; pass `chain_type=Any` to get that. =#
+function AbstractMCMC.sample(
+    rng::Random.AbstractRNG,
+    model::DensityModel,
+    sampler::Union{MALASampler,AdaptiveMALASampler,ParallelMALASampler},
+    N_or_isdone;
+    chain_type::Type=VNChain,
+    kwargs...,
+)
+    return AbstractMCMC.mcmcsample(rng, model, sampler, N_or_isdone; chain_type, kwargs...)
+end
+
+function AbstractMCMC.sample(
+    rng::Random.AbstractRNG,
+    model::DensityModel,
+    sampler::Union{MALASampler,AdaptiveMALASampler,ParallelMALASampler},
+    parallel::AbstractMCMC.AbstractMCMCEnsemble,
+    N::Integer,
+    nchains::Integer;
+    chain_type::Type=VNChain,
+    kwargs...,
+)
+    return AbstractMCMC.mcmcsample(
+        rng, model, sampler, parallel, N, nchains; chain_type, kwargs...
+    )
+end
diff --git a/test/test-AbstractMCMC-Interface.jl b/test/test-AbstractMCMC-Interface.jl
index b520c51..c274312 100644
--- a/test/test-AbstractMCMC-Interface.jl
+++ b/test/test-AbstractMCMC-Interface.jl
@@ -131,7 +131,8 @@ gradlogp_iface(x) = -x
         sampler = MALASampler(0.1)
 
         chain = sample(model, sampler, 100; progress=false)
-        @test length(chain) == 100
+        @test chain isa VNChain
+        @test FlexiChains.niters(chain) == 100
     end
 
     @testset "sample() with chain_type=Chains" begin
@@ -144,7 +145,7 @@ gradlogp_iface(x) = -x
         @test FlexiChains.niters(chain) == 200
 
         # Parameter columns present
-        @test only(FlexiChains.parameters(chain)) == @varname(x)
+        @test only(FlexiChains.parameters(chain)) == @vn(x)
 
         # Internal columns present
         @test Set(FlexiChains.extras(chain)) == Set(FlexiChains.Extra.([:logp, :accepted]))
@@ -226,7 +227,7 @@ gradlogp_iface(x) = -x
             )
 
             @test chain isa VNChain
-            @test FlexiChains.parameters(chain) == [@varname(mu), @varname(sigma)]
+            @test FlexiChains.parameters(chain) == [@vn(mu), @vn(sigma)]
         end
 
         @testset "with vector-valued symbols" begin
@@ -251,12 +252,12 @@ gradlogp_iface(x) = -x
                 50;
                 chain_type=VNChain,
                 progress=false,
-                param_names=[@varname(param) => (2,)],
+                param_names=[@vn(param) => (2,)],
             )
 
             @test chain isa VNChain
-            @test FlexiChains.parameters(chain) == [@varname(param)]
-            @test size(chain[@varname(param), stack = true]) == (50, 1, 2)
+            @test FlexiChains.parameters(chain) == [@vn(param)]
+            @test size(chain[@vn(param), stack = true]) == (50, 1, 2)
         end
     end
 
@@ -287,7 +288,7 @@ gradlogp_iface(x) = -x
         @testset "VNChain" begin
             chain = sample(model, sampler, 50; chain_type=VNChain, progress=false)
             @test chain isa VNChain
-            @test FlexiChains.parameters(chain) == [@varname(mu), @varname(sigma)]
+            @test FlexiChains.parameters(chain) == [@vn(mu), @vn(sigma)]
         end
     end
 
diff --git a/test/test-Adaptive-MALA.jl b/test/test-Adaptive-MALA.jl
index fb68db8..85adccf 100644
--- a/test/test-Adaptive-MALA.jl
+++ b/test/test-Adaptive-MALA.jl
@@ -174,8 +174,9 @@ end
     model = DensityModel(logp_adapt, gradlogp_adapt, 2)
     sampler = AdaptiveMALASampler(0.2; n_warmup=50)
 
-    samples = sample(MersenneTwister(1), model, sampler, 100; progress=false)
-    @test length(samples) == 100
+    chain = sample(MersenneTwister(1), model, sampler, 100; progress=false)
+    @test chain isa VNChain
+    @test FlexiChains.niters(chain) == 100
 end
 
 @testset "sample() with chain_type=Chains" begin
diff --git a/test/test-DEER-Interface.jl b/test/test-DEER-Interface.jl
index 130ae65..1a04862 100644
--- a/test/test-DEER-Interface.jl
+++ b/test/test-DEER-Interface.jl
@@ -205,8 +205,9 @@ end
     model = DensityModel(logp_deer, gradlogp_deer, 2)
     sampler = ParallelMALASampler(0.05; T=16, backend=_AD)
 
-    samples = sample(MersenneTwister(1), model, sampler, 50; progress=false)
-    @test length(samples) == 50
+    chain = sample(MersenneTwister(1), model, sampler, 50; progress=false)
+    @test chain isa VNChain
+    @test FlexiChains.niters(chain) == 50
 end
 
 @testset "ParallelMALASampler sample() with chain_type=Chains" begin
@@ -291,7 +292,14 @@ end
         @test fast[:logp] == slow[:logp]
 
         raw = sample(
-            MersenneTwister(7), model, sampler, N; initial_params=x0, progress=false, kw...
+            MersenneTwister(7),
+            model,
+            sampler,
+            N;
+            chain_type=Any,
+            initial_params=x0,
+            progress=false,
+            kw...,
         )
         slow_raw = ParallelMCMC._default_parallel_mala_mcmcsample(
             MersenneTwister(7), model, sampler, N; initial_params=x0, progress=false, kw...
diff --git a/test/test-DEER-Logistic.jl b/test/test-DEER-Logistic.jl
index 939294b..730ab5e 100644
--- a/test/test-DEER-Logistic.jl
+++ b/test/test-DEER-Logistic.jl
@@ -198,7 +198,14 @@ else
 
         n_samples, n_burn = 400, 100
 
-        raw_cpu = sample(MersenneTwister(42), model_cpu, sampler, n_samples; progress=false)
+        raw_cpu = sample(
+            MersenneTwister(42),
+            model_cpu,
+            sampler,
+            n_samples;
+            chain_type=Any,
+            progress=false,
+        )
         β_cpu = vec(mean(reduce(hcat, [s.x for s in raw_cpu[(n_burn + 1):end]]); dims=2))
 
         raw_gpu = sample(
@@ -206,6 +213,7 @@ else
             model_gpu,
             sampler,
             n_samples;
+            chain_type=Any,
             initial_params=CUDA.CuArray(zeros(Float32, _LR_D)),
             progress=false,
         )
diff --git a/test/test-GPU-AD-HVP.jl b/test/test-GPU-AD-HVP.jl
index e841ec2..2aa8cff 100644
--- a/test/test-GPU-AD-HVP.jl
+++ b/test/test-GPU-AD-HVP.jl
@@ -99,6 +99,7 @@ else
             model,
             sampler,
             n_samples;
+            chain_type=Any,
             initial_params=x0,
             progress=false,
         )
@@ -206,6 +207,7 @@ else
             model,
             sampler,
             n_samples;
+            chain_type=Any,
             initial_params=x0,
             progress=false,
         )
@@ -256,6 +258,7 @@ else
             model,
             sampler,
             n_samples;
+            chain_type=Any,
             initial_params=x0,
             progress=false,
         )
diff --git a/test/test-GPU-DEER.jl b/test/test-GPU-DEER.jl
index 2636815..eb8e2fd 100644
--- a/test/test-GPU-DEER.jl
+++ b/test/test-GPU-DEER.jl
@@ -213,6 +213,7 @@ else
             model,
             sampler,
             3_000;
+            chain_type=Any,
             initial_params=CUDA.CuArray(zeros(Float32, D)),
             progress=false,
         )
diff --git a/test/test-GPU-Performance.jl b/test/test-GPU-Performance.jl
index 1e99343..2e1459d 100644
--- a/test/test-GPU-Performance.jl
+++ b/test/test-GPU-Performance.jl
@@ -79,6 +79,7 @@ else
             model,
             sampler,
             sampler isa ParallelMALASampler ? sampler.T : 100;
+            chain_type=Any,
             progress=false,
             initial_params=x0,
         )
@@ -95,7 +96,13 @@ else
             end
             t0 = time_ns()
             sample(
-                MersenneTwister(42), model, sampler, N; progress=false, initial_params=x0
+                MersenneTwister(42),
+                model,
+                sampler,
+                N;
+                chain_type=Any,
+                progress=false,
+                initial_params=x0,
             )
             if x0 isa CUDA.CuArray
                 CUDA.synchronize()
diff --git a/test/test-LogDensityProblems.jl b/test/test-LogDensityProblems.jl
index 89d9542..97b82c0 100644
--- a/test/test-LogDensityProblems.jl
+++ b/test/test-LogDensityProblems.jl
@@ -55,8 +55,8 @@ end
         discard_warmup=true,
         progress=false,
     )
-    @test only(FlexiChains.parameters(chain)) == @varname(μ)
-    draws = vec(chain[@varname(μ)])
+    @test only(FlexiChains.parameters(chain)) == @vn(μ)
+    draws = vec(chain[@vn(μ)])
     @test abs(mean(draws) - LDP_MU_POST) < 0.1
     @test abs(var(draws) - LDP_VAR_POST) < 0.05
 end