Merge pull request #20 from fsaad/main

Fix bug in JAX array initialization
AaltoML · May 2, 2023 · f7b1744 · f7b1744
2 parents ad56794 + aa52e19
commit f7b1744
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/bayesnewton/ops.py b/bayesnewton/ops.py
@@ -237,7 +237,7 @@ def parallel_filtering_operator(elem1, elem2):
 
 
 def make_associative_filtering_elements(As, Qs, H, ys, noise_covs, m0, P0):
-    Qs = Qs.set(0).set(P0)  # first element requires different initialisation
+    Qs = Qs.at[0].set(P0)  # first element requires different initialisation
     AA, b, C, J, eta = parallel_filtering_element(As, Qs, H, noise_covs, ys)
     # modify initial b to account for m0 (not needed if m0=zeros)
     S = H @ Qs[0] @ H.T + noise_covs[0]