top_sparse3 now computes all elements sparsely

tests/test_utils.py

@@ -45,11 +45,8 @@ def test_top_sparse3(self):
         top_d = np.einsum(einstr, *mode_factors)
 
         # Now get the top numerator for the sparse method
-        mode_factors = [a for j, a in enumerate(factors) if j != factor]
         top_s = np.zeros(factors[factor].shape, dtype=np.float32)
-        utils.top_sparse3(x_indices, x_vals, top_s, beta, factor, model,
-                          mode_factors[0], mode_factors[1])
-
+        utils.top_sparse3(x_indices, x_vals, top_s, beta, factor, *factors)
 
         result = np.allclose(top_d, top_s, rtol=1e-5)
         self.assertTrue(result)

utils.py

@@ -4,7 +4,7 @@
 alphabet = 'abcdefghijklmnopqrstuvwxyz'
 
 
-def top_sparse3(x_indices, x_vals, out, beta, factor, model, A, B):
+def top_sparse3(x_indices, x_vals, out, beta, factor, A, B, C):
     # In einstein notation with factor=0 this is 'bz,cz,abc->az'
     # In einstein notation with factor=1 this is 'az,cz,abc->bz'
     # In einstein notation with factor=2 this is 'az,bz,abc->cz'
@@ -13,15 +13,18 @@ def top_sparse3(x_indices, x_vals, out, beta, factor, model, A, B):
     assert factor in (0, 1, 2), "Factor index must be < rank"
     if factor == 0:
         for (a, b, c), val in zip(x_indices, x_vals):
-            temp = val * A[b, :] * B[c, :] * (model[a, b, c] ** (beta - 2))
+            core = np.sum(A[a, :] * B[b, :] * C[c, :])
+            temp = val * B[b, :] * C[c, :] * (core ** (beta - 2))
             out[a, :] += temp
     if factor == 1:
         for (a, b, c), val in zip(x_indices, x_vals):
-            temp = val * A[a, :] * B[c, :] * (model[a, b, c] ** (beta - 2))
+            core = np.sum(A[a, :] * B[b, :] * C[c, :])
+            temp = val * A[a, :] * C[c, :] * (core ** (beta - 2))
             out[b, :] += temp
     elif factor == 2:
         for (a, b, c), val in zip(x_indices, x_vals):
-            temp = val * A[a, :] * B[b, :] * (model[a, b, c] ** (beta - 2))
+            core = np.sum(A[a, :] * B[b, :] * C[c, :])
+            temp = val * A[a, :] * B[b, :] * (core ** (beta - 2))
             out[c, :] += temp