fix: inconsistent return dims (#558)

* fix: inconsistent return dims

* test: inconsistent indexing

* fix: inconsistent dimensions inside gather getindex

Author: avik-pal

src/TracedRArray.jl

@@ -65,6 +65,9 @@ function Base.getindex(
 end
 
 Base.getindex(a::TracedRArray{T,0}) where {T} = TracedRNumber{T}((), a.mlir_data)
+function Base.getindex(a::TracedRArray{T,0}, ::CartesianIndex{0}) where {T}
+    return TracedRNumber{T}((), a.mlir_data)
+end
 
 function generate_index_list(i1, is...)
     list = reshape(i1, :, 1) .- 1
@@ -123,7 +126,11 @@ function Base.getindex(a::TracedRArray{T,N}, indices::CartesianIndex{N}) where {
     indices =
         materialize_traced_array(
             reshape(
-                TracedUtils.promote_to(TracedRArray{Int,1}, vcat(Tuple(indices)...)), 1, N
+                TracedUtils.promote_to(
+                    TracedRArray{Int,1}, collect(Int64, vcat(Tuple(indices)...))
+                ),
+                1,
+                N,
             ),
         ) .- 1
     return Ops.gather_getindex(a, indices)[1]
@@ -157,10 +164,19 @@ function Base.getindex(a::TracedRArray{T,N}, indices::Vararg{Any,N}) where {T,N}
         if any(i -> unwrapped_eltype(i) <: Bool, indices)
             error("Boolean indexing with TracedRArrays isn't fully supported yet.")
         end
-        indices_list = map(Base.Fix1(TracedUtils.promote_to, TracedRArray{Int,1}), indices)
+        idxs = map(indices) do i
+            i isa Number && return fill(i, 1)
+            return i
+        end
+        indices_list = map(Base.Fix1(TracedUtils.promote_to, TracedRArray{Int,1}), idxs)
         indices_list = generate_index_list(indices_list...)
         res = Ops.gather_getindex(a, indices_list)
-        return Ops.reshape(res, length.(indices)...)
+        res = Ops.reshape(res, length.(idxs)...)
+        ddims = findall(indices) do idx
+            return idx isa Integer || idx isa TracedRNumber{<:Integer}
+        end
+        isempty(ddims) || return materialize_traced_array(dropdims(res; dims=Tuple(ddims)))
+        return res
     end
 
     start_indices = map(indices) do i
@@ -172,7 +188,9 @@ function Base.getindex(a::TracedRArray{T,N}, indices::Vararg{Any,N}) where {T,N}
     )
 
     x = TracedRArray{T,N}((), res, Tuple(length.(indices)))
-    ddims = findall(Base.Fix2(isa, Integer), indices)
+    ddims = findall(indices) do idx
+        return idx isa Integer || idx isa TracedRNumber{<:Integer}
+    end
     isempty(ddims) || return materialize_traced_array(dropdims(x; dims=Tuple(ddims)))
     return x
 end
@@ -306,6 +324,11 @@ TracedUtils.promote_to(::Type{TracedRArray{T,N}}, rhs) where {T,N} = TracedRArra
 function TracedUtils.promote_to(::TracedRArray{T,N}, rhs) where {T,N}
     return TracedUtils.promote_to(TracedRArray{T,N}, rhs)
 end
+function TracedUtils.promote_to(
+    ::Type{TracedRArray{T,0}}, rhs::TracedRNumber{T2}
+) where {T,T2}
+    return TracedRArray{T,0}((), Ops.convert(TracedRNumber{T}, rhs).mlir_data, ())
+end
 
 for (jlop, hloop, hlocomp, merge) in
     ((:(Base.:(==)), :compare, "EQ", :all), (:(Base.:(!=)), :compare, "NE", :any))

src/TracedUtils.jl

@@ -374,7 +374,7 @@ end
 function elem_apply(f, args::Vararg{Any,Nargs}) where {Nargs}
     if all(iszero ∘ ndims, args)
         scalar_args = map(args) do arg
-            return promote_to(TracedRNumber{eltype(arg)}, arg)
+            return promote_to(TracedRNumber{Reactant.unwrapped_eltype(arg)}, arg)
         end
         return f(scalar_args...)
     end

test/basic.jl

@@ -1051,3 +1051,16 @@ end
     @test isfinite(Reactant.to_rarray(0.0; track_numbers=Number))
     @test !isfinite(Reactant.to_rarray(Inf; track_numbers=Number))
 end
+
+@testset "inconsistent indexing" begin
+    x_ra = Reactant.to_rarray(rand(3, 4, 3))
+    idx_ra = Reactant.to_rarray(1; track_numbers=Number)
+
+    fn1(x) = x[:, :, 1]
+    fn2(x, idx) = x[:, :, idx]
+    fn3(x, idx) = x[idx, :, 1]
+
+    @test ndims(@jit(fn1(x_ra))) == 2
+    @test ndims(@jit(fn2(x_ra, idx_ra))) == 2
+    @test ndims(@jit(fn3(x_ra, idx_ra))) == 1
+end