Merge pull request #1242 from bogiebro/master

Fix #1241

Author: ToucheSir

src/lib/base.jl

@@ -29,11 +29,14 @@ function accum(a::AbstractDict, b::AbstractDict)
 end
 
 @adjoint function getindex(d::AbstractDict, k)
-  d[k], function (Δ)
+  val = d[k]
+  function dict_getindex_pullback(Δ)
+    accum_param(__context__, val, Δ) === nothing && return
     grad = grad_mut(__context__, d)
     grad[k] = accum(get(grad, k, nothing), Δ)
     return (grad, nothing)
   end
+  val, dict_getindex_pullback
 end
 
 @adjoint! function setindex!(d::AbstractDict, v, k)

test/lib/base.jl

@@ -0,0 +1,13 @@
+@testset "base.jl" begin
+    @testset "Dict getindex with implicit params" begin
+        d = Dict{String, Vector{Float64}}("key"=>ones(4))
+        fn() = d["key"][2]
+        result1 = gradient(fn, Params([d["key"]]))[d["key"]]
+
+        x = d["key"]
+        fn2() = x[2]
+        result2 = gradient(fn2, Params([x]))[x]
+
+        @test result1 == result2
+    end
+end

test/runtests.jl

@@ -29,6 +29,7 @@ using CUDA: has_cuda
   @testset "lib" begin
     include("lib/number.jl")
     include("lib/lib.jl")
+    include("lib/base.jl")
     include("lib/array.jl")
   end