Argmax over columns of a large matrix

General Usage Performance
1

Hi there,

I have a large matrix for e.g (7x10000) and I want to iterate over each of the columns and compute the index of the maximum element. There are naive ways I can think of but since this part of my code is performance critical I wanted to ask more experienced users what would be the most efficient approach to do so?

Thanks a lot in advance!

2

I’d keep it simple with the builtin argmax. There are two ways I would try, the dimensional argmax and map/broadcast over eachcol.

julia> argmax(rand(7,5); dims=1)
1×5 Matrix{CartesianIndex{2}}:
 CartesianIndex(6, 1)  CartesianIndex(2, 2)  CartesianIndex(2, 3)  CartesianIndex(3, 4)  CartesianIndex(7, 5)

julia> map(argmax, eachcol(rand(7,5)))
5-element Vector{Int64}:
 4
 7
 6
 3
 3

When I benchmark these on your 7x10000 matrix, I find the map solution to be faster.

Although `argmin()` is much slower than `findmin()` due to automatic inline · Issue #56375 · JuliaLang/julia · GitHub suggests that argmax might be slower than necessary at present, I expect that will be fixed soon™ so would prefer to stick to builtins. In a few minutes of trying, I couldn’t write a version that was faster for your length=7 columns anyway, so that issue might be irrelevant at this size.

1 Like
3

Thanks a lot!