From 53dbcc7ed2af011df26dca3a1de7917c8039ca5e Mon Sep 17 00:00:00 2001 From: Anish Mehta Date: Wed, 14 Oct 2026 12:45:49 +0530 Subject: [PATCH] Fix bfmi with non-default chain_axis and draw_axis --- src/arviz_stats/base/array.py | 2 +- tests/test_minimal.py | 13 +++++++++++++ 2 files changed, 14 insertions(+), 1 deletion(-) diff --git a/src/arviz_stats/base/array.py b/src/arviz_stats/base/array.py index 85eab8e2..240e46bf 100644 --- a/src/arviz_stats/base/array.py +++ b/src/arviz_stats/base/array.py @@ -374,7 +374,7 @@ def bfmi(self, ary, chain_axis=-2, draw_axis=-1): n_dims=2, ravel=False, ) - return bfmi_ufunc(ary, out_shape=(ary.shape[chain_axis],)) + return bfmi_ufunc(ary, out_shape=(ary.shape[-2],)) def pareto_khat( self, ary, chain_axis=-2, draw_axis=-1, r_eff=None, tail="both", log_weights=False diff --git a/tests/test_minimal.py b/tests/test_minimal.py index 4c18d137..2a5b5aaf 100644 --- a/tests/test_minimal.py +++ b/tests/test_minimal.py @@ -190,6 +190,19 @@ def test_bfmi_datatree_returns_datatree(): assert result0.shape == (4,) +@pytest.mark.parametrize("axis", ["01", "20"]) +def test_bfmi_axes(axis, data_c0d1, data_c2d0): + if axis == "01": + bfmi = array_stats.bfmi(data_c0d1, chain_axis=0, draw_axis=1) + energy = np.moveaxis(data_c0d1, (0, 1), (-2, -1)) + else: + bfmi = array_stats.bfmi(data_c2d0, chain_axis=2, draw_axis=0) + energy = np.moveaxis(data_c2d0, (2, 0), (-2, -1)) + expected = np.square(np.diff(energy, axis=-1)).mean(axis=-1) / np.var(energy, axis=-1, ddof=1) + assert bfmi.shape == (3, 4) + np.testing.assert_allclose(bfmi, expected) + + @pytest.mark.parametrize("axis", ["01", "20"]) def test_pareto_min_ss_axes(axis, data_c0d1, data_c2d0): if axis == "01":