Fix getindex (#202)

Author: devmotion

Project.toml

@@ -1,6 +1,6 @@
 name = "ReverseDiff"
 uuid = "37e2e3b7-166d-5795-8a7a-e32c996b4267"
-version = "1.14.0"
+version = "1.14.1"
 
 [deps]
 ChainRulesCore = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4"

src/tracked.jl

@@ -280,8 +280,6 @@ Base.promote_rule(::Type{TrackedReal{V1,D1,O1}}, ::Type{TrackedReal{V2,D2,O2}})
 # AbstractArray Interface #
 ###########################
 
-Base.getindex(t::TrackedArray, i::Int) = TrackedReal(value(t)[i], deriv(t)[i], tape(t), i, t)
-
 colon2range(s, i) = i
 colon2range(s, ::Colon) = s
 
@@ -296,10 +294,10 @@ function index_iterable(shape::NTuple{N,Any}, i::NTuple{M,Any}) where {N,M}
 end
 
 for T in (:AbstractRange, :Colon, :(Union{Colon,AbstractRange}))
-    @eval function Base.getindex(t::TrackedArray, i::$(T)...)
+    @eval Base.@propagate_inbounds function Base.getindex(t::TrackedArray, i1::$(T), is::$(T)...)
         tp = tape(t)
-        out = TrackedArray(value(t)[i...], deriv(t)[i...], tp)
-        idx = index_iterable(axes(t), i)
+        out = TrackedArray(value(t)[i1, is...], deriv(t)[i1, is...], tp)
+        idx = index_iterable(axes(t), (i1, is...))
         record!(tp, SpecialInstruction, getindex, (t, idx), out)
         return out
     end
@@ -329,24 +327,25 @@ end
     return nothing
 end
 
-function Base.getindex(t::TrackedArray, inds::AbstractArray{<:CartesianIndex})
+Base.@propagate_inbounds function Base.getindex(t::TrackedArray, inds::AbstractArray{<:CartesianIndex})
     tp = tape(t)
     out = TrackedArray(value(t)[inds], deriv(t)[inds], tp)
     record!(tp, SpecialInstruction, getindex, (t, inds), out)
     return out
 end
-function Base.getindex(t::TrackedArray, i::Int...)
-    ind = LinearIndices(t)[i...]
-    return TrackedReal(value(t)[i...], deriv(t)[i...], tape(t), ind, t)
+Base.@propagate_inbounds function Base.getindex(t::TrackedArray, i1::Integer, is::Integer...)
+    ind = LinearIndices(t)[i1, is...]
+    return TrackedReal(value(t)[i1, is...], deriv(t)[i1, is...], tape(t), ind, t)
 end
-function Base.getindex(t::TrackedArray, _inds::Union{Integer, Colon, AbstractArray{<:Integer}}...)
-    inds = ntuple(Val(length(_inds))) do i
-        _inds[i] isa Colon && return firstindex(t,i):lastindex(t,i)
-        return _inds[i]
+Base.@propagate_inbounds function Base.getindex(t::TrackedArray, _inds1::Union{Integer, Colon, AbstractArray{<:Integer}}, _inds2::Union{Integer, Colon, AbstractArray{<:Integer}}...)
+    inds1 = _inds1 isa Colon ? axes(t, 1) : _inds1
+    inds2 = ntuple(Val(length(_inds2))) do i
+        _inds2[i] isa Colon && return axes(t, i+1)
+        return _inds2[i]
     end
     tp = tape(t)
-    out = TrackedArray(value(t)[inds...], deriv(t)[inds...], tp)
-    record!(tp, SpecialInstruction, (getindex, Val(:generic)), (t, inds), out)
+    out = TrackedArray(value(t)[inds1, inds2...], deriv(t)[inds1, inds2...], tp)
+    record!(tp, SpecialInstruction, (getindex, Val(:generic)), (t, (inds1, inds2...)), out)
     return out
 end
 @noinline function special_reverse_exec!(instruction::SpecialInstruction{<:Tuple{typeof(getindex), Val{:generic}}})

test/TrackedTests.jl

@@ -616,7 +616,12 @@ ta = TrackedArray(varr, darr, tp)
 
 @test isa(similar(ta), Matrix{eltype(ta)})
 
-@test samefields(ta[2], TrackedReal(varr[2], darr[2], tp, 2, ta))
+for T in (UInt, Int)
+    @test samefields(ta[T(2)], TrackedReal(varr[2], darr[2], tp, 2, ta))
+    @test samefields(ta[T(2), T(3)], TrackedReal(varr[2, 3], darr[2, 3], tp, 8, ta))
+    S = T === UInt ? Int : UInt
+    @test samefields(ta[S(2), T(3)], TrackedReal(varr[2, 3], darr[2, 3], tp, 8, ta))
+end
 
 ta_sub = ta[:,:]
 idx = ReverseDiff.index_iterable(axes(ta), (:, :))
@@ -630,53 +635,58 @@ instr = tp[1]
 @test instr.cache === nothing
 empty!(tp)
 
-ta_sub = ta[:,1:2]
-idx = ReverseDiff.index_iterable(axes(ta), (:, 1:2))
-@test collect(idx) == [(i, j) for i in 1:3, j in 1:2]
-@test samefields(ta_sub, TrackedArray(varr[:,1:2], darr[:,1:2], tp))
-@test length(tp) == 1
-instr = tp[1]
-@test instr.func === getindex
-@test instr.input === (ta, idx)
-@test samefields(instr.output, TrackedArray(varr[:,1:2], darr[:,1:2], tp))
-@test instr.cache === nothing
-empty!(tp)
-
-ta_sub = ta[2:3,:]
-idx = ReverseDiff.index_iterable(axes(ta), (2:3, :))
-@test collect(idx) == [(i, j) for i in 2:3, j in 1:3]
-@test samefields(ta_sub, TrackedArray(varr[2:3,:], darr[2:3,:], tp))
-@test length(tp) == 1
-instr = tp[1]
-@test instr.func === getindex
-@test instr.input === (ta, idx)
-@test samefields(instr.output, TrackedArray(varr[2:3,:], darr[2:3,:], tp))
-@test instr.cache === nothing
-empty!(tp)
-
-ta_sub = ta[1:2,2:3]
-idx = ReverseDiff.index_iterable(axes(ta), (1:2, 2:3))
-@test collect(idx) == [(i, j) for i in 1:2, j in 2:3]
-@test samefields(ta_sub, TrackedArray(varr[1:2,2:3], darr[1:2,2:3], tp))
-@test length(tp) == 1
-instr = tp[1]
-@test instr.func === getindex
-@test instr.input === (ta, idx)
-@test samefields(instr.output, TrackedArray(varr[1:2,2:3], darr[1:2,2:3], tp))
-@test instr.cache === nothing
-empty!(tp)
+for T in (UInt, Int)
+    ta_sub = ta[:,T(1):T(2)]
+    idx = ReverseDiff.index_iterable(axes(ta), (:, T(1):T(2)))
+    @test collect(idx) == [(i, j) for i in 1:3, j in 1:2]
+    @test samefields(ta_sub, TrackedArray(varr[:,1:2], darr[:,1:2], tp))
+    @test length(tp) == 1
+    instr = tp[1]
+    @test instr.func === getindex
+    @test instr.input === (ta, idx)
+    @test samefields(instr.output, TrackedArray(varr[:,1:2], darr[:,1:2], tp))
+    @test instr.cache === nothing
+    empty!(tp)
+
+    ta_sub = ta[T(2):T(3),:]
+    idx = ReverseDiff.index_iterable(axes(ta), (T(2):T(3), :))
+    @test collect(idx) == [(i, j) for i in 2:3, j in 1:3]
+    @test samefields(ta_sub, TrackedArray(varr[2:3,:], darr[2:3,:], tp))
+    @test length(tp) == 1
+    instr = tp[1]
+    @test instr.func === getindex
+    @test instr.input === (ta, idx)
+    @test samefields(instr.output, TrackedArray(varr[2:3,:], darr[2:3,:], tp))
+    @test instr.cache === nothing
+    empty!(tp)
+
+    S = T === UInt ? Int : UInt
+    for U in (S, T)
+        ta_sub = ta[S(1):S(2),T(2):T(3)]
+        idx = ReverseDiff.index_iterable(axes(ta), (S(1):S(2), T(2):T(3)))
+        @test collect(idx) == [(i, j) for i in 1:2, j in 2:3]
+        @test samefields(ta_sub, TrackedArray(varr[1:2,2:3], darr[1:2,2:3], tp))
+        @test length(tp) == 1
+        instr = tp[1]
+        @test instr.func === getindex
+        @test instr.input === (ta, idx)
+        @test samefields(instr.output, TrackedArray(varr[1:2,2:3], darr[1:2,2:3], tp))
+        @test instr.cache === nothing
+        empty!(tp)
+    end
 
-ta_sub = ta[2:6]
-idx = ReverseDiff.index_iterable(axes(ta), (2:6,))
-@test collect(idx) == [(i,) for i in 2:6]
-@test samefields(ta_sub, TrackedArray(varr[2:6], darr[2:6], tp))
-@test length(tp) == 1
-instr = tp[1]
-@test instr.func === getindex
-@test instr.input === (ta, idx)
-@test samefields(instr.output, TrackedArray(varr[2:6], darr[2:6], tp))
-@test instr.cache === nothing
-empty!(tp)
+    ta_sub = ta[T(2):T(6)]
+    idx = ReverseDiff.index_iterable(axes(ta), (T(2):T(6),))
+    @test collect(idx) == [(i,) for i in 2:6]
+    @test samefields(ta_sub, TrackedArray(varr[2:6], darr[2:6], tp))
+    @test length(tp) == 1
+    instr = tp[1]
+    @test instr.func === getindex
+    @test instr.input === (ta, idx)
+    @test samefields(instr.output, TrackedArray(varr[2:6], darr[2:6], tp))
+    @test instr.cache === nothing
+    empty!(tp)
+end
 
 ta_sub = ta[:]
 idx = ReverseDiff.index_iterable(axes(ta), (:,))