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