missing `Colon` dispatch

[yolhan83]

Hello, I come accross this quite a lot so I thought I would make an mwe here

function foo(x,i)
      return x[1:i]
 end
x = Reactant.to_rarray(ones(10));
i = Reactant.to_rarray(2;track_numbers=true);
@compile foo(x,i)

error :

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

Closest candidates are:
  (::Colon)(::T, ::Any, ::T) where T<:Real
   @ Base range.jl:50
  (::Colon)(::A, ::Any, ::C) where {A<:Real, C<:Real}
   @ Base range.jl:10
  (::Colon)(::T, ::Any, ::T) where T
   @ Base range.jl:49
  ...

Stacktrace:
  [1] macro expansion
    @ ~/.julia/packages/Reactant/bR5J7/src/utils.jl:0 [inlined]
  [2] call_with_reactant(::Reactant.MustThrowError, ::Colon, ::Int64, ::Reactant.TracedRNumber{Int64})
    @ Reactant ~/.julia/packages/Reactant/bR5J7/src/utils.jl:893
  [3] foo
    @ ./REPL[36]:2 [inlined]
  [4] (::Nothing)(none::typeof(foo), none::Reactant.TracedRArray{Float64, 1}, none::Reactant.TracedRNumber{Int64})
    @ Reactant ./<missing>:0
  [5] foo
    @ ./REPL[36]:2 [inlined]
  [6] call_with_reactant(::typeof(foo), ::Reactant.TracedRArray{Float64, 1}, ::Reactant.TracedRNumber{Int64})
    @ Reactant ~/.julia/packages/Reactant/bR5J7/src/utils.jl:0
  [7] make_mlir_fn(f::typeof(foo), 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/bR5J7/src/TracedUtils.jl:348
  [8] 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/bR5J7/src/Compiler.jl:1594
  [9] compile_mlir! (repeats 2 times)
    @ ~/.julia/packages/Reactant/bR5J7/src/Compiler.jl:1561 [inlined]
 [10] compile_xla(f::Function, args::Tuple{…}; before_xla_optimizations::Bool, client::Nothing, serializable::Bool, kwargs::@Kwargs{…})
    @ Reactant.Compiler ~/.julia/packages/Reactant/bR5J7/src/Compiler.jl:3483
 [11] compile_xla
    @ ~/.julia/packages/Reactant/bR5J7/src/Compiler.jl:3456 [inlined]
 [12] compile(f::Function, args::Tuple{…}; kwargs::@Kwargs{…})
    @ Reactant.Compiler ~/.julia/packages/Reactant/bR5J7/src/Compiler.jl:3555
 [13] top-level scope
    @ ~/.julia/packages/Reactant/bR5J7/src/Compiler.jl:2633
Some type information was truncated. Use `show(err)` to see complete types.

[wsmoses]

want to take a first attempt at adding the method (presumably just promoting the first to a traced and calling the other)

[yolhan83]

I think for now we don't have a method

function foo(x,i)
    return x[Reactant.TracedRNumber{Int}(1):i]
end

ERROR: LoadError: TypeError: non-boolean (Reactant.TracedRNumber{Bool}) used in boolean context
Stacktrace:
  [1] Colon
    @ ./range.jl:7 [inlined]
  [2] (::Nothing)(none::Colon, none::Reactant.TracedRNumber{Int64}, none::Reactant.TracedRNumber{Int64})
    @ Reactant ./<missing>:0
  [3] string
    @ ./strings/io.jl:189 [inlined]
  [4] macro expansion
    @ ~/.julia/packages/Reactant/bR5J7/src/Ops.jl:50 [inlined]
  [5] >=
    @ ~/.julia/packages/Reactant/bR5J7/src/TracedRNumber.jl:298 [inlined]
  [6] Colon
    @ ./range.jl:7 [inlined]
  [7] call_with_reactant(::Reactant.MustThrowError, ::Colon, ::Reactant.TracedRNumber{Int64}, ::Reactant.TracedRNumber{Int64})
    @ Reactant ~/.julia/packages/Reactant/bR5J7/src/utils.jl:0
  [8] foo
    @ /mnt/c/Users/yolha/Desktop/juju_tests/port_ml/main.jl:4 [inlined]
  [9] (::Nothing)(none::typeof(foo), none::Reactant.TracedRArray{Float64, 1}, none::Reactant.TracedRNumber{Int64})
    @ Reactant ./<missing>:0
 [10] TracedRNumber
    @ ~/.julia/packages/Reactant/bR5J7/src/TracedRNumber.jl:169 [inlined]
 [11] foo
    @ /mnt/c/Users/yolha/Desktop/juju_tests/port_ml/main.jl:4 [inlined]
 [12] call_with_reactant(::typeof(foo), ::Reactant.TracedRArray{Float64, 1}, ::Reactant.TracedRNumber{Int64})
    @ Reactant ~/.julia/packages/Reactant/bR5J7/src/utils.jl:0
 [13] make_mlir_fn(f::typeof(foo), 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/bR5J7/src/TracedUtils.jl:348
 [14] 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/bR5J7/src/Compiler.jl:1594
 [15] compile_mlir! (repeats 2 times)
    @ ~/.julia/packages/Reactant/bR5J7/src/Compiler.jl:1561 [inlined]
 [16] compile_xla(f::Function, args::Tuple{ConcretePJRTArray{…}, ConcretePJRTNumber{…}}; before_xla_optimizations::Bool, client::Nothing, serializable::Bool, kwargs::@Kwargs{compile_options::CompileOptions, fn_kwargs::@NamedTuple{}})
    @ Reactant.Compiler ~/.julia/packages/Reactant/bR5J7/src/Compiler.jl:3483
 [17] compile_xla
    @ ~/.julia/packages/Reactant/bR5J7/src/Compiler.jl:3456 [inlined]
 [18] compile(f::Function, args::Tuple{…}; kwargs::@Kwargs{…})
    @ Reactant.Compiler ~/.julia/packages/Reactant/bR5J7/src/Compiler.jl:3555
 [19] top-level scope
    @ ~/.julia/packages/Reactant/bR5J7/src/Compiler.jl:2633
 [20] include(fname::String)
    @ Main ./sysimg.jl:38
 [21] top-level scope
    @ REPL[5]:1
in expression starting at /mnt/c/Users/yolha/Desktop/juju_tests/port_ml/main.jl:10
Some type information was truncated. Use `show(err)` to see complete types.

do we have a way to make a traced iterable without allocating the array ?

[avik-pal]

Reactant.jl/src/Types.jl

Lines 99 to 105 in 1ec3452

	struct TracedUnitRange{T} <: AbstractUnitRange{T}
	start::T
	stop::T
	function TracedUnitRange{T}(start::T, stop::T) where {T}
	return new(start, unitrange_last(start, stop))
	end
	end

[yolhan83]

Ok so this seems enough for colon

function (C::Base.Colon)(start::Reactant.TracedRNumber{T},stop::Reactant.TracedRNumber{T}) where T
    Reactant.TracedUnitRange(start,stop)
end
function (C::Base.Colon)(start::Reactant.TracedRNumber{T},stop::T) where T
    C(start,Reactant.TracedRNumber{T}(stop))
end
function (C::Base.Colon)(start::T,stop::Reactant.TracedRNumber{T}) where T
    C(Reactant.TracedRNumber{T}(start),stop)
end

with step I tried with step and using TracedStepRangeLen but I get tons of show errors but the compilation works.

However, indexing an array into it is hard I'm trying to avoid an opcall by going back to a method defined for AbstractArrays using slices (ie getindex_linear) but with not much results for now.

[avik-pal]

Let's get this part merged.

Indexing an array with this is a bit more non-trivial, since we need to setup dynamic sized array support for that.

[yolhan83]

I think the rest of this is a bit too early to think about we should close ?

[yolhan83]

Ps : if you need anything on Julia side don't hesitate, I might one day learn c++ but I won't even try to touch Enzyme-JAX before I feel ready