change stacked ordering

cf TuringLang/MCMCDiagnosticTools.jl#50
tpapp · Dec 13, 2022 · eebe8f5 · eebe8f5
1 parent 9253328
commit eebe8f5
Show file tree

Hide file tree

Showing 3 changed files with 7 additions and 3 deletions.
diff --git a/Project.toml b/Project.toml
@@ -22,6 +22,7 @@ DocStringExtensions = "0.8, 0.9"
 FillArrays = "0.13"
 LogDensityProblems = "1, 2"
 LogExpFunctions = "0.3"
+MCMCDiagnosticTools = "0.2"
 Parameters = "0.11, 0.12"
 ProgressMeter = "1"
 TensorCast = "0.4"

diff --git a/src/mcmc.jl b/src/mcmc.jl
@@ -561,12 +561,15 @@ $(SIGNATURES)
 
 Given a vector of `results`, each containing a property `posterior_matrix` (eg obtained from
 [`mcmc_with_warmup`](@ref) with the same sample length), return a lazy view as an array
-indexed by `[draw_index, parameter_index, chain_index]`.
+indexed by `[draw_index, chain_index, parameter_index]`.
 
 This is useful as an input for eg `MCMCDiagnosticTools.ess_rhat`.
+
+!!! note
+    The ordering is not compatible with MCMCDiagnostictools version < 0.2.
 """
 function stack_posterior_matrices(results)
-    @cast _[i, j, k]:= results[k].posterior_matrix[j, i]
+    @cast _[i, k, j]:= results[k].posterior_matrix[j, i]
 end
 
 """

diff --git a/test/test_mcmc.jl b/test/test_mcmc.jl
@@ -74,7 +74,7 @@ end
     D, N, K = 5, 100, 7
     ℓ = multivariate_normal(ones(5))
     results = fill(mcmc_with_warmup(RNG, ℓ, N; reporter = NoProgressReport()), K)
-    @test size(stack_posterior_matrices(results)) == (N, D, K)
+    @test size(stack_posterior_matrices(results)) == (N, K, D)
     @test size(pool_posterior_matrices(results)) == (D, N * K)
 end