From b68cd8b6561846ec9f0e8326bf98a69f8ab07920 Mon Sep 17 00:00:00 2001 From: "Feras A. Saad" Date: Tue, 2 May 2023 18:43:57 -0400 Subject: [PATCH] Fix bug in _parallel_kf_mf.build_block_diag, use block_index instead of 0. Using Pzeros.at[0] gives a dimension error. All other instances of "Pzeros.at" in this module use Pzeros.at[block_index] instead of Pzeros.at[0]. --- bayesnewton/ops.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/bayesnewton/ops.py b/bayesnewton/ops.py index 747dadd..155cc2a 100644 --- a/bayesnewton/ops.py +++ b/bayesnewton/ops.py @@ -569,7 +569,7 @@ def _parallel_kf_mf(As, Qs, H, ys, noise_covs, m0, P0, masks, block_index): @vmap def build_block_diag(P_blocks): - P = Pzeros.at[0].add(P_blocks.flatten()) + P = Pzeros.at[block_index].add(P_blocks.flatten()) return P def build_mean(m):