Skip to content

Commit

Permalink
Recursive version that tries to choose a good order
Browse files Browse the repository at this point in the history
  • Loading branch information
dahong67 committed Mar 1, 2024
1 parent 21bebdb commit a071fbd
Showing 1 changed file with 9 additions and 7 deletions.
16 changes: 9 additions & 7 deletions src/gcp-opt.jl
Original file line number Diff line number Diff line change
Expand Up @@ -248,12 +248,14 @@ function khatrirao(A::Vararg{T,N}) where {T<:AbstractMatrix,N}
return A[1]
end

# General case: N > 1
r = (only unique)(size.(A, 2))
K = similar(A[1], prod(size.(A, 1)), r)
for j in 1:r
temp = (N == 2) ? view(A[1], :, j) : kron([view(A[i], :, j) for i in 1:N-1]...)
kron!(view(K, :, j), temp, view(A[N], :, j))
# Base case: N = 2
if N == 2
r = (only unique)(size.(A, 2))
return reshape(reshape(A[1], :, 1, r) .* reshape(A[2], 1, :, r), :, r)
end
return K

# Recursive case: N > 2
I, r = size.(A, 1), (only unique)(size.(A, 2))
n = argmin(n -> I[n] * I[n+1], 1:N-1)
return khatrirao(A[1:n-1]..., khatrirao(A[n], A[n+1]), A[n+2:end]...)
end

0 comments on commit a071fbd

Please sign in to comment.