diff --git a/bayesnewton/ops.py b/bayesnewton/ops.py
index 9523086..747dadd 100644
--- 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]