Broadcasting a 2D array with a multidimensional array in Julia

Question

I have a 2D matrix, A, of size m x n, where n could be a very large number (e.g. n > 10000) and a multidimensional matrix of size m x m x n. So, for every column i in A, I want to compute A[:,i]'*B[:,:,i]. Below are the codes that I tried with the broadcasting feature in Julia. However, the performance of my code is quite slow. I am wondering whether the performance of my code could be improved. So, does someone have an idea of how I can improve the codes?

using LinearAlgebra;
m = 500;
n = 20000; # this could be a very large number.

vecA = rand(m,n);
matB = rand(m,m,n);


combinedAB = Array{Array{Float64,2},2}(undef,n,1);
for ii in eachindex(combinedAB)
  combinedAB[ii] = [vecA[:,ii] matB[:,:,ii]];
end

# this is the result.
res = broadcast(eAB -> dotProd(eAB), combinedAB);

function dotProd(matZ::Array{Float64,2})
   return sum(broadcast(dot,matZ[:,1],matZ[:,2:end]),dims=1);
end

Bogumił Kamiński Bogumił Kamiński · Accepted Answer · 2019-06-14T12:44:35

Is this fast enough in your case?

res = [a'*b for (a, b) in zip(eachcol(vecA), eachslice(matB, dims=3))]

I do not have enough RAM to test is against your input values, but given the tests I did for smaller data it should run in ~3 seconds.

I also assume you really want an adjoint of a (this is what you have written in your question; if you are working with reals it should not matter if you use ' or transpose)

The key difference between the codes (it is hidden under the hood in my solution as additionally it is shorter) is that my solution does not allocate intermediate arrays but uses views.

Broadcasting a 2D array with a multidimensional array in Julia

1 Answers