# Broadcasting matrix-vector multiplication over specific axis for n-dimensional arrays

I have a matrix-vector multiplication to fill `axis=2` of my 3D array `a` (where `basis` is type `Matrix`):

``````@inbounds for n ∈ 1:size(a, 1), s ∈ 1:size(a, 3)
a[n, :, s] = @views basis * b[n, :, s]
end
``````

The question is how can I broadcast this over `axis=3`, so that the loop only goes through `n ∈ 1:size(a, 1)`.
Is there a performance benefit on doing so?

Thanks!

There’s nothing wrong with writing a loop. Done properly, it is usually at least as fast as broadcasting. However, you’ll get better performance from in-place multiplication `mul!` here (since you already have the output space allocated). Also, as you were perhaps alluding to, there’s no need to loop over `s` as that part can be handled by using matrix-matrix multiplication instead of matrix-vector.

``````using LinearAlgebra
for n in axes(a,1)
@views mul!(a[n,:,:], basis, b[n,:,:])
end
``````

I removed `@inbounds` because there isn’t a check here to ensure that `axes(a,1) == axes(b,1)` and also because it is unlikely to significantly change performance.

You can write the above via a somewhat arcane broadcast:

``````using LinearAlgebra
mul!.(eachslice(a;dims=1), Ref(basis), eachslice(b;dims=1))
``````

Although I wouldn’t say this is easier to read that the loop and I don’t expect it to perform noticeably differently. Note that I have broadcasted `mul!` by using `mul!.`, sliced the arrays by iterating the first dimension, and used `Ref(basis)` to not broadcast over that.

2 Likes

Yeap, of course the matrix-matrix multiplication would do this already… my bad.
And thanks for the insightful answer and the performance tips using `mul!`

Tangentially related, but this function may be useful to you. It is for selecting 1D slices in an ND array along a dimension.

1 Like