Fix Independent kernel

Fixes a minor bug that was introduced when fixing the deprecated jax.ops.index_update and jax.ops.index_add.
AaltoML · Jun 14, 2023 · a3a0589 · a3a0589
1 parent 8ce9598
commit a3a0589
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/bayesnewton/kernels.py b/bayesnewton/kernels.py
@@ -1586,11 +1586,11 @@ def __init__(self, kernels):
     def K(self, X, X2):
         zeros = np.zeros(self.num_kernels)
         K0 = self.kernel0.K(X, X2)
-        index_vector = index_vector.at[0].add(1.)
+        index_vector = zeros.at[0].add(1.)
         Kstack = np.kron(K0, np.diag(index_vector))
         for i in range(1, self.num_kernels):
             kerneli = eval("self.kernel" + str(i))
-            index_vector = index_vector.at[i].add(1.)
+            index_vector = zeros.at[i].add(1.)
             Kstack += np.kron(kerneli.K(X, X2), np.diag(index_vector))
         return Kstack