Does not work Reactant arrays

[Moelf]

julia> using ArraysOfArrays, Reactant

julia> cpu_ary=VectorOfVectors([rand(Float32, rand(2:5)) for _ = 1:3])
3-element VectorOfVectors{Float32, Vector{Float32}, Vector{Int64}, Vector{Tuple{}}}:
 Float32[0.46416235, 0.41684252, 0.08079648, 0.0020227432, 0.36202437]
 Float32[0.44927096, 0.12560415]
 Float32[0.18007994, 0.3705178, 0.69251794]

julia> r_ary = Reactant.to_rarray(cpu_ary);

julia> f = @compile map(argmin, r_ary);
ERROR: TypeError: in VectorOfArrays, in VT, expected VT<:AbstractVector{Float32}, got Type{Reactant.TracedRArray{Float32, 1}}
Stacktrace:
  [1] traced_type_inner(T::Type, seen::Dict{Type, Type}, mode::Reactant.TraceMode, track_numbers::Type, sharding::Any, runtime::Any)
    @ Reactant ~/.julia/packages/Reactant/wAaVF/src/Tracing.jl:734
  [2] traced_type(T::Type, ::Val{Reactant.ConcreteToTraced}, track_numbers::Type, sharding::Reactant.Sharding.NoSharding, runtime::Val{:PJRT})
    @ Reactant ~/.julia/packages/Reactant/wAaVF/src/Tracing.jl:870
  [3] make_tracer_unknown(seen::Reactant.OrderedIdDict{Any, Any}, prev::Any, path::Any, mode::Reactant.TraceMode; track_numbers::Type, sharding::Any, runtime::Any, kwargs::@Kwargs{toscalar::Bool})
    @ Reactant ~/.julia/packages/Reactant/wAaVF/src/Tracing.jl:1032
  [4] make_tracer_unknown
    @ ~/.julia/packages/Reactant/wAaVF/src/Tracing.jl:1009 [inlined]
  [5] #make_tracer#128
    @ ~/.julia/packages/Reactant/wAaVF/src/Tracing.jl:1146 [inlined]
  [6] make_tracer
    @ ~/.julia/packages/Reactant/wAaVF/src/Tracing.jl:1136 [inlined]
  [7] prepare_mlir_fn_args(args::Tuple{…}, name::String, concretein::Bool, toscalar::Bool, argprefix::Symbol, runtime::Val{…}, optimize_then_pad::Bool, do_transpose::Bool, input_shardings::Nothing, verify_arg_names::Nothing)
    @ Reactant.TracedUtils ~/.julia/packages/Reactant/wAaVF/src/TracedUtils.jl:434
  [8] make_mlir_fn(f::typeof(map), args::Tuple{…}, kwargs::@NamedTuple{}, 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, num_replicas::Int64, optimize_then_pad::Bool)
    @ Reactant.TracedUtils ~/.julia/packages/Reactant/wAaVF/src/TracedUtils.jl:306
  [9] compile_mlir!(mod::Reactant.MLIR.IR.Module, f::Function, args::Tuple{…}, compile_options::CompileOptions, callcache::Dict{…}, sdycache::Dict{…}; fn_kwargs::@NamedTuple{}, backend::String, runtime::Val{…}, legalize_stablehlo_to_mhlo::Bool, kwargs::@Kwargs{})
    @ Reactant.Compiler ~/.julia/packages/Reactant/wAaVF/src/Compiler.jl:1538
 [10] compile_mlir! (repeats 2 times)
    @ ~/.julia/packages/Reactant/wAaVF/src/Compiler.jl:1505 [inlined]
 [11] compile_xla(f::Function, args::Tuple{…}; before_xla_optimizations::Bool, client::Nothing, serializable::Bool, kwargs::@Kwargs{…})
    @ Reactant.Compiler ~/.julia/packages/Reactant/wAaVF/src/Compiler.jl:3414
 [12] compile_xla
    @ ~/.julia/packages/Reactant/wAaVF/src/Compiler.jl:3387 [inlined]
 [13] compile(f::Function, args::Tuple{…}; kwargs::@Kwargs{…})
    @ Reactant.Compiler ~/.julia/packages/Reactant/wAaVF/src/Compiler.jl:3486
 [14] top-level scope
    @ ~/.julia/packages/Reactant/wAaVF/src/Compiler.jl:2567
Some type information was truncated. Use `show(err)` to see complete types.

[oschulz]

I'm not sure if this can work, with VectorOfVectors. From what I know, Reactant compiles a function for specific data sizes. We'll have to ask the experts if "ragged arrays" are possible with the current machinery.

[oschulz]

However, ArraysOfSimilarArrays should in principle be Reactant-compatible. At the moment @jit map(cumsum, AOSA) doesn't work, but it does work for Base.ColumnSlices. I think the problem is that ArraysOfArrays currently uses Vector{...} and so on as element type instead of SubArray{...}. This will require a breaking release of ArraysOfArrays, but it's probably unavoidable.

This means finally adressing #2.

Base.Slices uses Base.promote_op to figure out the correct element type (which works for CuArray too, it seems, where a view is not a SubArray but another CuArray):

function Slices(A::P, slicemap::SM, ax::AX) where {P,SM,AX}
    N = length(ax)
    argT = map((a,l) -> l === (:) ? Colon : eltype(a), axes(A), slicemap)
    S = Base.promote_op(view, P, argT...)
    Slices{P,SM,AX,S,N}(A, slicemap, ax)
end

So this should work for us, too.

[Moelf]

idk, I think the problem is Reactant may need an Ext (like this https://github.com/EnzymeAD/Reactant.jl/blob/main/ext/ReactantOffsetArraysExt.jl) for every Julia package that implements some kind of Array container...

[oschulz]

I think the problem is Reactant may need an Ext

Yes, could be. Also, I looked into @jit map(cumsum, AOSA) for eachslice again - it works, but the result are actually separate Reactant arrays, not views into a single array, so that's also not what should be possible yet.

[oschulz]

But as a first step, it's probably time to change eltype from Array to the actual view type. We can do a v1.0 release.

[oschulz]

Having accumulated a bit more experience with Reactant now, yes, this will require a Reactant ArraysOfArrays extensions for ArrayOfSimilarArrays. And VectorOfArrays may not be Reactant-capable for general operations because XLA/StableHLO wants fixed shapes.