bug in comprehension over dictionaries

[rlrs]

It currently seems quite difficult to use Dicts with Zygote - at least I have trouble getting it to work at all. Here's a MWE where the Dict is obviously unnecessary but it illustrates my point nonetheless. I could be doing something wrong, however this is a common data structure used in a straightforward fashion.

using Zygote
w = randn(5)
function f(w)
    d = Dict{Int, Float64}(i => v for (i,v) in enumerate(w))
    sum([v for (_, v) in d])
end
gradient(f, w)

Error: map is not defined on dictionaries.
Using iterator syntax for the sum instead of the list comprehension syntax gives: MethodError: no method matching getindex(::Dict{Any,Any}), which probably makes the issue similar to #408.

[rlrs]

I've been digging a bit, and the issue (for the list comprehension syntax) seems to be this definition:

Zygote.jl/src/lib/base.jl

Line 25 in 469e34b

grad_mut(d::AbstractDict) = Dict()

The reason is that in literal_getproperty:

Zygote.jl/src/lib/lib.jl

Lines 199 to 212 in 469e34b

	@adjoint function literal_getproperty(x, ::Val{f}) where f
	val = getproperty(x, f)
	function back(Δ)
	accum_param(__context__, val, Δ) === nothing && return
	if isimmutable(x)
	((;nt_nothing(x)...,pair(Val(f), Δ)...), nothing)
	else
	dx = grad_mut(__context__, x)
	dx[] = (;dx[]...,pair(Val(f),accum(getfield(dx[], f), Δ))...)
	return (dx,nothing)
	end
	end
	unwrap(val), back
	end

the expected grad_mut seems to be this one:

Zygote.jl/src/lib/lib.jl

Line 226 in 469e34b

grad_mut(x) = Ref{Any}(nt_nothing(x))

which returns a Ref to a NamedTuple with field names of the struct:

julia> Ref{Any}(Zygote.nt_nothing(Dict()))
Base.RefValue{Any}((slots = nothing, keys = nothing, vals = nothing, ndel = nothing, count = nothing, age = nothing, idxfloor = nothing, maxprobe = nothing))

Of course, I expect there would be further errors after fixing this. I don't have a good enough overview of the code to make a fix, so any guidance would be appreciated.

[CarloLucibello]

Other examples with stacktraces

julia> d = Dict(1 => 5.0, 2 => 6.0);

julia> gradient(d -> sum([v for (_,v) in d]), d)
ERROR: map is not defined on dictionaries
Stacktrace:
  [1] error(s::String)
    @ Base ./error.jl:33
  [2] map(f::Function, #unused#::Dict{Int64, Float64})
    @ Base ./abstractarray.jl:2325
  [3] ∇map(cx::Zygote.Context, f::var"#12#14", args::Dict{Int64, Float64})
    @ Zygote ~/.julia/dev/Zygote/src/lib/array.jl:203
  [4] _pullback(cx::Zygote.Context, #unused#::typeof(collect), g::Base.Generator{Dict{Int64, Float64}, var"#12#14"})
    @ Zygote ~/.julia/dev/Zygote/src/lib/array.jl:249
  [5] _pullback
    @ ./REPL[8]:1 [inlined]
  [6] _pullback(ctx::Zygote.Context, f::var"#11#13", args::Dict{Int64, Float64})
    @ Zygote ~/.julia/dev/Zygote/src/compiler/interface2.jl:0
  [7] _pullback(f::Function, args::Dict{Int64, Float64})
    @ Zygote ~/.julia/dev/Zygote/src/compiler/interface.jl:34
  [8] pullback(f::Function, args::Dict{Int64, Float64})
    @ Zygote ~/.julia/dev/Zygote/src/compiler/interface.jl:40
  [9] gradient(f::Function, args::Dict{Int64, Float64})
    @ Zygote ~/.julia/dev/Zygote/src/compiler/interface.jl:75
 [10] top-level scope
    @ REPL[8]:1

and

julia> d = Dict(1 => 5.0, 2 => 6.0);

julia> gradient(d -> sum(v for (_,v) in d), d)
ERROR: MethodError: no method matching getindex(::Dict{Any, Any})
Closest candidates are:
  getindex(::Dict{K, V}, ::Any) where {K, V} at dict.jl:480
  getindex(::AbstractDict, ::Any) at abstractdict.jl:494
  getindex(::AbstractDict, ::Any, ::Any, ::Any...) at abstractdict.jl:504
Stacktrace:
  [1] (::Zygote.var"#back#222"{:vals, Zygote.Context, Dict{Int64, Float64}, Vector{Float64}})(Δ::Zygote.OneElement{Float64, 1, Tuple{Int64}, Tuple{Base.OneTo{Int64}}})
    @ Zygote ~/.julia/dev/Zygote/src/lib/lib.jl:233
  [2] (::Zygote.var"#1789#back#223"{Zygote.var"#back#222"{:vals, Zygote.Context, Dict{Int64, Float64}, Vector{Float64}}})(Δ::Zygote.OneElement{Float64, 1, Tuple{Int64}, Tuple{Base.OneTo{Int64}}})
    @ Zygote ~/.julia/packages/ZygoteRules/OjfTt/src/adjoint.jl:59
  [3] Pullback
    @ ./Base.jl:33 [inlined]
  [4] (::typeof(∂(getproperty)))(Δ::Zygote.OneElement{Float64, 1, Tuple{Int64}, Tuple{Base.OneTo{Int64}}})
    @ Zygote ~/.julia/dev/Zygote/src/compiler/interface2.jl:0
  [5] Pullback
    @ ~/.julia/packages/ZygoteRules/OjfTt/src/ZygoteRules.jl:11 [inlined]
  [6] (::typeof(∂(literal_getproperty)))(Δ::Zygote.OneElement{Float64, 1, Tuple{Int64}, Tuple{Base.OneTo{Int64}}})
    @ Zygote ~/.julia/dev/Zygote/src/compiler/interface2.jl:0
  [7] Pullback
    @ ./dict.jl:684 [inlined]
  [8] Pullback
    @ ./dict.jl:688 [inlined]
  [9] (::typeof(∂(iterate)))(Δ::Tuple{NamedTuple{(:first, :second), Tuple{Nothing, Float64}}, Nothing})
    @ Zygote ~/.julia/dev/Zygote/src/compiler/interface2.jl:0
 [10] Pullback
    @ ./reduce.jl:60 [inlined]
 [11] (::typeof(∂(_foldl_impl)))(Δ::Float64)
    @ Zygote ~/.julia/dev/Zygote/src/compiler/interface2.jl:0
 [12] Pullback
    @ ./reduce.jl:48 [inlined]
 [13] (::typeof(∂(foldl_impl)))(Δ::Float64)
    @ Zygote ~/.julia/dev/Zygote/src/compiler/interface2.jl:0
 [14] Pullback
    @ ./reduce.jl:44 [inlined]
 [15] (::typeof(∂(mapfoldl_impl)))(Δ::Float64)
    @ Zygote ~/.julia/dev/Zygote/src/compiler/interface2.jl:0
 [16] Pullback (repeats 2 times)
    @ ./reduce.jl:160 [inlined]
 [17] (::typeof(∂(mapfoldl)))(Δ::Float64)
    @ Zygote ~/.julia/dev/Zygote/src/compiler/interface2.jl:0
 [18] Pullback
    @ ./reduce.jl:287 [inlined]
 [19] (::typeof(∂(#mapreduce#218)))(Δ::Float64)
    @ Zygote ~/.julia/dev/Zygote/src/compiler/interface2.jl:0
 [20] Pullback
    @ ./reduce.jl:287 [inlined]
 [21] (::typeof(∂(mapreduce)))(Δ::Float64)
    @ Zygote ~/.julia/dev/Zygote/src/compiler/interface2.jl:0
 [22] Pullback
    @ ./reduce.jl:501 [inlined]
 [23] (::typeof(∂(#sum#221)))(Δ::Float64)
    @ Zygote ~/.julia/dev/Zygote/src/compiler/interface2.jl:0
 [24] Pullback
    @ ./reduce.jl:501 [inlined]
 [25] (::typeof(∂(sum)))(Δ::Float64)
    @ Zygote ~/.julia/dev/Zygote/src/compiler/interface2.jl:0
 [26] Pullback
    @ ./reduce.jl:528 [inlined]
 [27] (::typeof(∂(#sum#222)))(Δ::Float64)
    @ Zygote ~/.julia/dev/Zygote/src/compiler/interface2.jl:0
 [28] Pullback
    @ ./reduce.jl:528 [inlined]
 [29] (::typeof(∂(sum)))(Δ::Float64)
    @ Zygote ~/.julia/dev/Zygote/src/compiler/interface2.jl:0
 [30] Pullback
    @ ./REPL[7]:1 [inlined]
 [31] (::typeof(∂(#7)))(Δ::Float64)
    @ Zygote ~/.julia/dev/Zygote/src/compiler/interface2.jl:0
 [32] (::Zygote.var"#50#51"{typeof(∂(#7))})(Δ::Float64)
    @ Zygote ~/.julia/dev/Zygote/src/compiler/interface.jl:41
 [33] gradient(f::Function, args::Dict{Int64, Float64})
    @ Zygote ~/.julia/dev/Zygote/src/compiler/interface.jl:76
 [34] top-level scope
    @ REPL[7]:1

[CarloLucibello]

Generator syntax works now

julia> d = Dict(1 => 5.0, 2 => 6.0);

julia> gradient(d -> sum(v for (_,v) in d), d)
(Dict{Any, Any}(2 => 1.0, 1 => 1.0),)

Comprehension [..] still gives the same error