Fix bug in JAX array initialization: Qs.set(0).set(P0) -> Qs.at[0].se…

…t(P0).

The current code fails with the following AttributeError:

    --> 240     Qs = Qs.set(0).set(P0)  # first element requires different initialisation
    AttributeError: DynamicJaxprTracer has no attribute set

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]