Refactor Jn_inner and Kn_inner computation

dahong67 · Jan 9, 2024 · 2a65e1a · 2a65e1a
1 parent 18e5213
commit 2a65e1a
Showing 1 changed file with 2 additions and 5 deletions.
diff --git a/src/gcp-opt.jl b/src/gcp-opt.jl
@@ -205,12 +205,9 @@ function mttkrp(X, U, n)
         kr_inner = khatrirao(U[reverse(n+1:N)]...) 
         kr_outer = khatrirao(U[reverse(1:n-1)]...)
         inner = reshape(reshape(X, Jn, Kn) * kr_inner, (size(X)[1:n]..., r)) 
+        Jn_inner = prod(size(inner)[1:n-1])
+        Kn_inner = prod(size(inner)[n:end-1])
         for j in 1:r
-            # Inner tensor-vector products
-            #inner = reshape(reshape(X, Jn, Kn) * kr_inner[:, j], size(X)[1:n]) 
-            # Outer tensor-vector products
-            Jn_inner = prod(size(selectdim(inner, ndims(inner), 1))[1:n-1])
-            Kn_inner = prod(size(selectdim(inner, ndims(inner), 1))[n:end])
             Rn[:, j] = transpose(reshape(selectdim(inner, ndims(inner), j), Jn_inner, Kn_inner)) * kr_outer[:, j]
         end
     end