mapreduce with multiple arrays? #1554
Description
mapreduce with multiple arrays currently doesn't seem to work:
julia> Reactant.set_default_backend("gpu")
julia> A = randn(Float32, 10^6); B = randn(Float32, 10^6);
julia> rA, rB = Reactant.to_rarray((A, B))
(ConcretePJRTArray{Float32, 1, 1, Reactant.Sharding.ShardInfo{Reactant.Sharding.NoSharding, Nothing}}(Float32[-0.018601555, 1.494709, 0.88219666, -0.5095448, 1.5401549, -0.5151698, -0.93188196, 0.9969609, -0.69384325, 0.031682458 … -0.12592782, 1.8223894, 1.078637, 0.9034769, 2.9029932, -1.6135824, 0.055411525, -0.44011563, -0.07745815, 0.4226902]), ConcretePJRTArray{Float32, 1, 1, Reactant.Sharding.ShardInfo{Reactant.Sharding.NoSharding, Nothing}}(Float32[-0.44233724, -1.248382, -0.35614863, 0.14472502, -0.7663462, 0.4234709, -0.89164025, -1.1270964, 0.06115119, -1.0243567 … 0.9129408, 0.7977728, -0.39093867, 0.047124397, -0.2396115, -0.37811318, -1.162326, 0.046417087, -1.3277433, 0.7805992]))
julia> foo(A, B) = mapreduce(+, +, A, B; init = zero(eltype(A)), dims = 1)
foo (generic function with 2 methods)
julia> foo(A, B)
1-element Vector{Float32}:
692.58276
julia> @jit foo(rA, rB)
ERROR: Scalar indexing is disallowed.
Invocation of getindex(::TracedRArray, ::Vararg{Int, N}) resulted in scalar indexing of a GPU array.
[...]In principle this should be possible, I guess (at some point)? Or is there a different "Reactant-way" to do this?