Argmax fails with reactant

[mkschleg]

In FluxML/OneHotArrays.jl#49 we found that reactant doesn't work with argmax on julia 1.11.4 in a fresh environment.

Package versions:

[7e8f7934] MLDataDevices v1.9.1
[3c362404] Reactant v0.2.59

MWE:

julia> using Reactant, MLDataDevices

julia> argmax([true false false; false true true]; dims=1)
1×3 Matrix{CartesianIndex{2}}:
 CartesianIndex(1, 1)  CartesianIndex(2, 2)  CartesianIndex(2, 3)

julia> argmax([true false false; false true true] |> reactant_device(); dims=1)
ERROR: MethodError: no method matching vec(::Tuple{Int64})

Full stack trace:

ERROR: MethodError: no method matching vec(::Tuple{Int64})
The function `vec` exists, but no method is defined for this combination of argument types.

Closest candidates are:
  vec(::SparseArrays.AbstractSparseVector)
   @ SparseArrays ~/.julia/juliaup/julia-1.11.4+0.aarch64.apple.darwin14/share/julia/stdlib/v1.11/SparseArrays/src/sparsevector.jl:1128
  vec(::LinearAlgebra.Adjoint{<:Real, <:AbstractVector})
   @ LinearAlgebra ~/.julia/juliaup/julia-1.11.4+0.aarch64.apple.darwin14/share/julia/stdlib/v1.11/LinearAlgebra/src/adjtrans.jl:374
  vec(::LinearAlgebra.Transpose{<:Number, <:AbstractVector})
   @ LinearAlgebra ~/.julia/juliaup/julia-1.11.4+0.aarch64.apple.darwin14/share/julia/stdlib/v1.11/LinearAlgebra/src/adjtrans.jl:373
  ...

Stacktrace:
  [1] macro expansion
    @ ~/.julia/packages/Reactant/AebXg/src/utils.jl:0 [inlined]
  [2] call_with_reactant(::typeof(vec), ::Tuple{Int64})
    @ Reactant ~/.julia/packages/Reactant/AebXg/src/utils.jl:790
  [3] getindex
    @ ~/.julia/packages/Reactant/AebXg/src/TracedRArray.jl:176 [inlined]
  [4] getindex(none::Reactant.TracedRArray{Bool, 2}, none::Tuple{Int64, CartesianIndex{1}})
    @ Reactant ./<missing>:0
  [5] getindex
    @ ~/.julia/packages/Reactant/AebXg/src/TracedRArray.jl:167 [inlined]
  [6] call_with_reactant(::typeof(getindex), ::Reactant.TracedRArray{Bool, 2}, ::Int64, ::CartesianIndex{1})
    @ Reactant ~/.julia/packages/Reactant/AebXg/src/utils.jl:0
  [7] make_mlir_fn(f::typeof(getindex), args::Tuple{…}, kwargs::Tuple{}, name::String, concretein::Bool; toscalar::Bool, return_dialect::Symbol, args_in_result::Symbol, construct_function_without_args::Bool, do_transpose::Bool, input_shardings::Nothing, output_shardings::Nothing, runtime::Val{…}, verify_arg_names::Nothing, argprefix::Symbol, resprefix::Symbol, resargprefix::Symbol)
    @ Reactant.TracedUtils ~/.julia/packages/Reactant/AebXg/src/TracedUtils.jl:303
  [8] make_mlir_fn
    @ ~/.julia/packages/Reactant/AebXg/src/TracedUtils.jl:178 [inlined]
  [9] compile_mlir!(mod::Reactant.MLIR.IR.Module, f::typeof(getindex), args::Tuple{…}, callcache::Dict{…}, sdycache::IdDict{…}; optimize::Bool, shardy_passes::Symbol, no_nan::Bool, backend::String, fn_kwargs::Tuple{}, raise::Bool, input_shardings::Nothing, output_shardings::Nothing, do_transpose::Bool, runtime::Val{…})
    @ Reactant.Compiler ~/.julia/packages/Reactant/AebXg/src/Compiler.jl:840
 [10] compile_mlir!
    @ ~/.julia/packages/Reactant/AebXg/src/Compiler.jl:784 [inlined]
 [11] compile_xla(f::Function, args::Tuple{ConcretePJRTArray{Bool, 2, 1, Reactant.Sharding.ShardInfo{Reactant.Sharding.NoSharding, Nothing}}, Int64, CartesianIndex{1}}; client::Nothing, kwargs::@Kwargs{})
    @ Reactant.Compiler ~/.julia/packages/Reactant/AebXg/src/Compiler.jl:1992
 [12] compile_xla
    @ ~/.julia/packages/Reactant/AebXg/src/Compiler.jl:1974 [inlined]
 [13] compile(f::Function, args::Tuple{ConcretePJRTArray{Bool, 2, 1, Reactant.Sharding.ShardInfo{Reactant.Sharding.NoSharding, Nothing}}, Int64, CartesianIndex{1}}; sync::Bool, kwargs::@Kwargs{})
    @ Reactant.Compiler ~/.julia/packages/Reactant/AebXg/src/Compiler.jl:2040
 [14] compile
    @ ~/.julia/packages/Reactant/AebXg/src/Compiler.jl:2039 [inlined]
 [15] getindex(::ConcretePJRTArray{Bool, 2, 1, Reactant.Sharding.ShardInfo{Reactant.Sharding.NoSharding, Nothing}}, ::Int64, ::CartesianIndex{1})
    @ Reactant ~/.julia/packages/Reactant/AebXg/src/ConcreteRArray.jl:271
 [16] findminmax!(f::typeof(identity), op::typeof(isless), Rval::ConcretePJRTArray{Bool, 2, 1, Reactant.Sharding.ShardInfo{…}}, Rind::Matrix{CartesianIndex{…}}, A::ConcretePJRTArray{Bool, 2, 1, Reactant.Sharding.ShardInfo{…}})
    @ Base ./reducedim.jl:1039
 [17] _findmax(f::typeof(identity), A::ConcretePJRTArray{Bool, 2, 1, Reactant.Sharding.ShardInfo{Reactant.Sharding.NoSharding, Nothing}}, region::Int64)
    @ Base ./reducedim.jl:1209
 [18] _findmax
    @ ./reducedim.jl:1176 [inlined]
 [19] findmax
    @ ./reducedim.jl:1175 [inlined]
 [20] #argmax#992
    @ ./reducedim.jl:1274 [inlined]
 [21] top-level scope
    @ REPL[5]:1
Some type information was truncated. Use `show(err)` to see complete types.

[avik-pal]

x_ra = Reactant.to_rarray([true false false; false true true])

@jit argmax(x_ra, dims=1)

you need to compile and run code for the most part.