From eebe8f5617bd0626c54b42611c6d7983345a8b1d Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Tam=C3=A1s=20K=2E=20Papp?= <tkpapp@gmail.com>
Date: Tue, 13 Dec 2022 10:19:25 +0100
Subject: [PATCH] change stacked ordering

cf https://github.com/TuringLang/MCMCDiagnosticTools.jl/pull/50
---
 Project.toml      | 1 +
 src/mcmc.jl       | 7 +++++--
 test/test_mcmc.jl | 2 +-
 3 files changed, 7 insertions(+), 3 deletions(-)

diff --git a/Project.toml b/Project.toml
index bee758e..75412a5 100644
--- 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
index 621a146..0c71169 100644
--- 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
index b274dd8..f640f31 100644
--- 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