Merge pull request #2 from FluxML/wct/fix-getindex

Fix indexing

Author: MikeInnes

src/lib/array.jl

@@ -107,6 +107,14 @@ Base.getindex(xs::TrackedArray, i...; kwargs...) = track(getindex, xs, i...; kwa
     end
 end
 
+@grad function getindex(xs::AbstractArray, i::Array...)
+  data(xs)[i...], function (Δ)
+    Δ′ = zero(xs)
+    @views Δ′[i...] .+= data(Δ)
+    (nobacksies(:getindex, Δ′), map(_->nothing, i)...)
+  end
+end
+
 Base.view(x::TrackedArray, inds...; kwargs...) = track(Base.view, x, inds...; kwargs...)
 
 @grad function view(x::AbstractArray, inds...; kwargs...)

test/tracker.jl

@@ -122,6 +122,11 @@ end
 
 end
 
+@testset "getindex (Nabla.jl - #139)" begin
+  z = [2, 3, 3]
+  @test gradtest(x->x[z], randn(MersenneTwister(123456), 3))
+end
+
 @test gradtest(x -> permutedims(x, [3,1,2]), rand(4,5,6))
 @test gradtest(x -> PermutedDimsArray(x, [3,1,2]), rand(4,5,6))