Reduce/Accumulate with Matrix Matrix Product

I would like to use the standard lib reduce and accumulate to calculate the matrix matrix products instead of the element wise products.

  mat = Array{ComplexF64}(undef, 2, 2, 3)
  mat[:, :, 1] = [1+0.5im 0.0+0.0im; 0.0+0.0im 1+0.5im]
  mat[:, :, 2] = [3+1.5im 2+0.0im; 0.0+1.0im 0.4+1im]
  mat[:, :, 3] = [0.4+1.4im 0.0+0.0im; 0.0+0.0im 0.4+1.4im]
mat_acc = accumulate((acc, x) -> x * acc, mat, dims=3, init=Matrix{ComplexF64}(I, 2, 2))
mat_acc = accumulate(*, mat, dims=3, init=Matrix{ComplexF64}(I, 2, 2))

neither the first nor the second method work like expected. If i would like the accumulation to be element wise i would have written something like

mat_acc = accumulate(.*, mat, dims=3, init=Matrix{ComplexF64}(I, 2, 2))

is there a way to use accumulate/reduce in this manner? Of course i know one could use loops to do this as well, but i think this is cleaner.

Here is some additional behaviour clarification

julia> isapprox(mat[:,:,1]*mat[:,:,2]*mat[:,:,3],reduce(*,mat,dims=3))
false

julia> isapprox(mat[:,:,1].*mat[:,:,2].*mat[:,:,3],reduce(*,mat,dims=3))
true

Hi, and welcome to the Julia community!

(Note that there is an incorrect newline in the mat definition, which makes copy-pasting it yield an error.)

Just to check, as it is not immediately fully clear to me what the output should be, if mat has size (2, 2, n), do you want as output mat[:, :, 1] * mat[:, :, 2] * ... * mat[:, :, n] (where the * are matrix-matrix products)?

EDIT: the above is the output I would imagine for reduce. For accumulate you would then want [mat[:, :, 1], mat[:, :, 1] * mat[:, :, 2], mat[:, :, 1] * mat[:, :, 2] * mat[:, :, 3], ...]? Do you want this as a Vector{Matrix} like I’ve written here, or a 3D Array (like mat itself)?

Hi you are right i thought that println would output a julia parsable multidimensional array, but it doesn’t. The matrix should be as you presumed the dimension should be (2,2,n) and you believe assumed the same behavior as I did. But it doesn’t perform a matrix multiplication, it performs an element wise multiplication. A 3 matrix would be preferred

You could use eachslice:

julia> mat = rand(2, 2, 3);

julia> isapprox.(reduce(*, eachslice(mat, dims=3), mat[:, :, 1] * mat[:, :, 2] * mat[:, :, 3]))
true

For accumulate(*, eachslice(mat, dims=3)) you will get a length-3 Vector of matrices. You could turn it into an Array{Float64, 3} by using cat and splatting, though this will not be the most performant solution.

julia> acc_prod = cat(accumulate(*, eachslice(mat, dims=3))..., dims=3);

julia> isapprox(acc_prod[:, :, 2], mat[:, :, 1] * mat[:, :, 2])
true

A faster option would be to accumulate in-place to slices of an output Array:

function accumulate_mmprod1(M)
    return cat(accumulate(*, eachslice(M, dims=3))..., dims=3)
end

function accumulate_mmprod2(M)
    out = similar(M)
    accumulate!(*, eachslice(out, dims=3), eachslice(M, dims=3))
    return out
end

using BenchmarkTools
M = rand(2, 2, 1000)
@btime accumulate_mmprod1($M);  # 16.191 ms (11027 allocations: 14.03 MiB)
@btime accumulate_mmprod2($M);  # 49.600 μs (1002 allocations: 125.03 KiB)

This is perfect. Thx. However, I am still convinced that the reduce and accumulate function should support the distinction between element wise operations and matrix operations for consistency’s sake.

I understand where you’re coming from, but I think the current implementation also makes sense. When reducing our 2 \times 2 \times 3 Array over the third dimension, in your mind you’re (implicitly) viewing the data as a length-3 Vector of matrices of size 2 \times 2. In this case the elements are indeed matrices, so * should refer to matrix multiplications and .* to elementwise multiplication. But equally, you can consider it as a 2 \times 2 Matrix with length-3 Vector elements. Reducing then means that you squash each of these Vectors to a single number. The operation we’re using here is then simple scalar multiplication *.

Both interpretations can make sense for me, it’s just that Julia went for the other convention. Using eachslice you can still explicitly indicate that you want to use the Vector-of-matrices interpretation. Personally, I would expect that
reduce(*, rand(10, 20, 30, 40), dims=(3, 4)) would result in a 10 \times 20 (\times 1 \times 1) Array, rather than in a DimensionMismatch error, so I do slightly prefer the current implementation.

By the way, note that each matrix product in the faster accumulate_mmprod2 still allocates. While this would stray from using reduce (at least in a somewhat straightforward fashion) you can get rid of these allocations using mul!:

using LinearAlgebra

function accumulate_mmprod3(M)
    out = similar(M)
    out[:, :, 1] .= @view(M[:, :, 1])
    for i = 2:size(M, 3)
        @views mul!(out[:, :, i], out[:, :, i-1], M[:, :, i])
    end
    return out
end

@btime accumulate_mmprod3($M);  # 26.900 μs (2 allocations: 31.31 KiB)

If 2x2 (or some small, fixed size, i.e., <10x10) is your actual use case, I would use a Vector{SMatrix} rather than Array{T,3} as your base storage.

julia> using StaticArrays

julia> mats = randn(SMatrix{2,2,ComplexF64}, 3)
3-element Vector{SMatrix{2, 2, ComplexF64}}:
 [-0.4392383048041226 + 0.2885524087340694im 0.2609678193778586 - 0.6283397627041533im; 0.26731221365964686 - 0.9657035571328186im -1.076195222622194 - 1.4320432097280107im]
 [-0.7480293767498994 - 1.439409630509693im 0.5252036010763754 - 0.06736677154409094im; -0.09778082269973562 - 0.4057247048468974im 0.04132403382775557 + 0.5277495173384432im]
 [0.3584778954143221 - 0.19720872393567213im 0.1947793911015316 + 0.888543168734062im; -0.3720518628295307 - 0.2533384249667726im 0.7992260358263662 + 0.4350418323460239im]

julia> prod(mats) # reduction with efficient SMatrix multiplications
2×2 SMatrix{2, 2, ComplexF64, 4} with indices SOneTo(2)×SOneTo(2):
 0.264904-0.100256im  -0.262841+0.775395im
 -1.14483+0.964583im  -0.084733-2.23622im

Note that with the slices being AbstractMatrix to begin with, and a simple Vector of those, prod is the only reduction function you need.