Mul! with 3-dimensional array

I am new to Julia. I need to multiply 3 dimensional array by 2 dimensional array. The simplified form is here:

vR = fill(1.0, (251, 51, 21));
Π = fill(1.0, (21, 21));

function mat_mul(A, B)
C = similar(A)
@inbounds for i in 1:51
C[:, i, :] = A[:, i, :] * B’
end
return C
end

EvR = mat_mul(vR, Π)

Is there a faster way? I wish something like `mul!`, but can be applied to 3-dimensional array.

This isn’t a full answer, but try to look at GitHub - mcabbott/Tullio.jl: ⅀ and other packages by mcabbott. They support the nifty Einstein notation (Einstein notation - Wikipedia)

2 Likes

NNlib.batched_mul! could also be used, you will need to reorder your axis though as it assumes that the batch dimensions are last.

2 Likes

The libraries people have suggested are probably your best bets if you really need top performance. But you probably get most of the way there with

``````using LinearAlgebra

function mat_mul(A, B)
C = similar(A)
for i in axes(A, 2)
@views mul!(C[:, i, :], A[:, i, :], B')
end
return C
end
``````

One can also just use

``````mat_mul2(A, B) = mapslices(a -> a * B', A; dims=(1,3))
``````

but this is rather inefficient due to the intermediate allocations.

1 Like