getindex assumes contiguous indexing

[avik-pal]

julia> using Reactant

julia> xx = rand(4, 4) |> Reactant.to_rarray
4×4 ConcreteRArray{Float64, 2}:
 0.333112  0.765684  0.808021  0.197128
 0.533312  0.582735  0.454427  0.9106
 0.323209  0.509113  0.918799  0.997749
 0.355191  0.437561  0.620295  0.506551

julia> foo(x) = x[[1, 3, 4], :]
foo (generic function with 1 method)

julia> @code_hlo foo(xx)
ERROR: AssertionError: size(MLIR.IR.type(mlir_data)) == shape
Stacktrace:
  [1] Reactant.TracedRArray{Float64, 2}(paths::Tuple{}, mlir_data::Reactant.MLIR.IR.Value, shape::Tuple{Int64, Int64})
    @ Reactant /mnt/software/lux/Reactant.jl/src/TracedRArray.jl:14
  [2] getindex(::Reactant.TracedRArray{Float64, 2}, ::Vector{Int64}, ::Function)
    @ Reactant /mnt/software/lux/Reactant.jl/src/TracedRArray.jl:103
  [3] foo
    @ ./REPL[17]:1 [inlined]
  [4] (::Tuple{})(none::Reactant.TracedRArray{Float64, 2})
    @ Base.Experimental ./<missing>:0
  [5] (::Reactant.var"#26#35"{Bool, typeof(foo), Vector{ConcreteRArray{Float64, 2}}, Vector{Union{ReactantCore.MissingTracedValue, Reactant.TracedRArray, Reactant.TracedRNumber}}, Tuple{Reactant.TracedRArray{Float64, 2}}})()
    @ Reactant /mnt/software/lux/Reactant.jl/src/utils.jl:148
  [6] block!(f::Reactant.var"#26#35"{Bool, typeof(foo), Vector{ConcreteRArray{Float64, 2}}, Vector{Union{ReactantCore.MissingTracedValue, Reactant.TracedRArray, Reactant.TracedRNumber}}, Tuple{Reactant.TracedRArray{Float64, 2}}}, blk::Reactant.MLIR.IR.Block)
    @ Reactant.MLIR.IR /mnt/software/lux/Reactant.jl/src/mlir/IR/Block.jl:201
  [7] make_mlir_fn(f::Function, args::Vector{ConcreteRArray{Float64, 2}}, kwargs::Tuple{}, name::String, concretein::Bool; toscalar::Bool, return_dialect::Symbol, no_args_in_result::Bool, construct_function_without_args::Bool)
    @ Reactant /mnt/software/lux/Reactant.jl/src/utils.jl:112
  [8] make_mlir_fn
    @ /mnt/software/lux/Reactant.jl/src/utils.jl:36 [inlined]
  [9] #6
    @ /mnt/software/lux/Reactant.jl/src/Compiler.jl:270 [inlined]
 [10] block!(f::Reactant.Compiler.var"#6#11"{typeof(foo), Vector{ConcreteRArray{Float64, 2}}}, blk::Reactant.MLIR.IR.Block)
    @ Reactant.MLIR.IR /mnt/software/lux/Reactant.jl/src/mlir/IR/Block.jl:201
 [11] #5
    @ /mnt/software/lux/Reactant.jl/src/Compiler.jl:269 [inlined]
 [12] mmodule!(f::Reactant.Compiler.var"#5#10"{Reactant.MLIR.IR.Module, typeof(foo), Vector{ConcreteRArray{Float64, 2}}}, blk::Reactant.MLIR.IR.Module)
    @ Reactant.MLIR.IR /mnt/software/lux/Reactant.jl/src/mlir/IR/Module.jl:93
 [13] compile_mlir!(mod::Reactant.MLIR.IR.Module, f::Function, args::Vector{ConcreteRArray{Float64, 2}}; optimize::Bool)
    @ Reactant.Compiler /mnt/software/lux/Reactant.jl/src/Compiler.jl:266
 [14] compile_mlir!
    @ /mnt/software/lux/Reactant.jl/src/Compiler.jl:265 [inlined]
 [15] #2
    @ /mnt/software/lux/Reactant.jl/src/Compiler.jl:260 [inlined]
 [16] context!(f::Reactant.Compiler.var"#2#3"{@Kwargs{optimize::Bool}, typeof(foo), Vector{ConcreteRArray{Float64, 2}}}, ctx::Reactant.MLIR.IR.Context)
    @ Reactant.MLIR.IR /mnt/software/lux/Reactant.jl/src/mlir/IR/Context.jl:76
 [17] compile_mlir(f::Function, args::Vector{ConcreteRArray{Float64, 2}}; kwargs::@Kwargs{optimize::Bool})
    @ Reactant.Compiler /mnt/software/lux/Reactant.jl/src/Compiler.jl:258
 [18] top-level scope
    @ /mnt/software/lux/Reactant.jl/src/Compiler.jl:419
 [19] top-level scope
    @ none:1

This is originating because we fix a stride of 1 and compute start and end with first and last

[avik-pal]

indices::Vararg{Union{Base.AbstractUnitRange,Colon,Int},N} --> tightening the signature to this, kind of solves the issue but the generated HLO is quite bad. We should be able to use gather to express this case I think