Add rule for Dict iteration

Author: ToucheSir

src/lib/base.jl

@@ -47,6 +47,29 @@ end
   end
 end
 
+# This rule behaves much like the getindex adjoint,
+# just with an (internal) ordinal index instead of a key.
+function _pullback(cx::AContext, ::typeof(iterate), d::Dict, i)
+  iter = iterate(d, i)
+  function dict_iterate_pullback(Δ)
+    (iter === nothing || Δ === nothing) && return
+    k, v = iter[1]
+    _, dv = Δ[1]
+    accum_param(cx, v, dv) === nothing && return
+    grad = grad_mut(cx, d)
+    grad[k] = accum(get(grad, k, nothing), dv)
+    return (nothing, grad, nothing)
+  end
+  return iter, dict_iterate_pullback
+end
+
+# ...while this one is to avoid duplicating code or differentiating skip_deleted.
+# The alternative would be to write a rule for the private _iterate(::Dict, i).
+function _pullback(cx::AContext, ::typeof(iterate), d::Dict)
+  # Calculation of i is the same used in iterate(::Dict)
+  return _pullback(cx, iterate, d, Base.skip_deleted(d, d.idxfloor))
+end
+
 # Channels
 
 grad_mut(ch::Channel) = Channel(ch.sz_max)

test/lib/base.jl

@@ -10,4 +10,23 @@
 
         @test result1 == result2
     end
+
+    @testset "Dict iteration" begin
+        # https://github.com/FluxML/Zygote.jl/issues/1065
+        function sumvals(d)
+            s = zero(d["c"])
+            for (k, v) in d
+                s += v
+                k == :b && (s += v)
+            end
+            return sum(s)
+        end
+
+        d = Dict(:a => 3, :b => 4, "c" => 5)
+        @test gradient(sumvals, d)[1] == Dict(:a => 1, :b => 2, "c" => 1)
+
+        d = Dict(:a => [3], :b => [4], "c" => [5])
+        grads = gradient(() -> sumvals(d), d |> values |> collect |> Params)
+        @test (grads[d[:a]], grads[d[:b]], grads[d["c"]]) == ([1], [2], [1])
+    end
 end