argmax of `A` over its first `N` dimensions and drops them. By default `N = 1`.

argmax_first.jl

"""
    argmax_first(A, Val(N) = Val(1))

argmax of `A` over its first `N` dimensions and drops them. By default `N = 1`.
"""
function argmax_first(A::AbstractArray, ::Val{N} = Val(1)) where {N}
    dims = tuplen(Val(N))
    argmax_(A; dims=dims)
end

@testset "argmax_first" begin
    @inferred argmax_first(randn(3,2,4))
    @test size(argmax_first(randn(3,2,4))) == (2,4)
    A = [3 4; 1 2]
    i = argmax_first(A)
    @test A[i] == [3, 4]
end