Bump FlexiChains compat to 0.7 and widen CUDA to 6

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.

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"

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.

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

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

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
 

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

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

test/test-DEER-Logistic.jl

@@ -198,14 +198,22 @@ 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(
             MersenneTwister(42),
             model_gpu,
             sampler,
             n_samples;
+            chain_type=Any,
             initial_params=CUDA.CuArray(zeros(Float32, _LR_D)),
             progress=false,
         )

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,
         )

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,
         )

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()

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