broken iteration over dictionary

[CarloLucibello]

Issue related to #725 :

julia> d = Dict(:a => 5, :b => 6);

julia> function f(d)
         s = 0
         for (k,v) in d
           s += v
         end
         s
       end
f (generic function with 1 method)

julia> gradient(f, 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{Symbol, Int64}, Vector{Int64}})(Δ::Zygote.OneElement{Int64, 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{Symbol, Int64}, Vector{Int64}}})(Δ::Zygote.OneElement{Int64, 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{Int64, 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{Int64, 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, Int64}}, Nothing})
    @ Zygote ~/.julia/dev/Zygote/src/compiler/interface2.jl:0
 [10] Pullback
    @ ./REPL[5]:4 [inlined]
 [11] (::typeof(∂(f)))(Δ::Int64)
    @ Zygote ~/.julia/dev/Zygote/src/compiler/interface2.jl:0
 [12] (::Zygote.var"#50#51"{typeof(∂(f))})(Δ::Int64)
    @ Zygote ~/.julia/dev/Zygote/src/compiler/interface.jl:41
 [13] gradient(f::Function, args::Dict{Symbol, Int64})
    @ Zygote ~/.julia/dev/Zygote/src/compiler/interface.jl:76
 [14] top-level scope
    @ REPL[6]:1

Also iteration using values is broken:

julia> function fval(d)
         s = 0
         for v in values(d)
           s += v
         end
         s
       end
fval (generic function with 1 method)

julia> gradient(fval, 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{Symbol, Int64}, Vector{Int64}})(Δ::Zygote.OneElement{Int64, 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{Symbol, Int64}, Vector{Int64}}})(Δ::Zygote.OneElement{Int64, 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{Int64, 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{Int64, 1, Tuple{Int64}, Tuple{Base.OneTo{Int64}}})
    @ Zygote ~/.julia/dev/Zygote/src/compiler/interface2.jl:0
  [7] Pullback
    @ ./dict.jl:697 [inlined]
  [8] (::typeof(∂(iterate)))(Δ::Tuple{Int64, Nothing})
    @ Zygote ~/.julia/dev/Zygote/src/compiler/interface2.jl:0
  [9] Pullback
    @ ./REPL[13]:4 [inlined]
 [10] (::typeof(∂(fval)))(Δ::Int64)
    @ Zygote ~/.julia/dev/Zygote/src/compiler/interface2.jl:0
 [11] (::Zygote.var"#50#51"{typeof(∂(fval))})(Δ::Int64)
    @ Zygote ~/.julia/dev/Zygote/src/compiler/interface.jl:41
 [12] gradient(f::Function, args::Dict{Symbol, Int64})
    @ Zygote ~/.julia/dev/Zygote/src/compiler/interface.jl:76
 [13] top-level scope
    @ REPL[14]:1

[mohdibntarek]

This showed up in JuliaNonconvex/Nonconvex.jl#130 so just giving it a bump.

[ToucheSir]

Dict handling is one of the most finicky bits of Zygote's internals and issues which involve the AD transform even more so. I'm not even sure any active contributors understand the latter in enough depth to easily root cause issues like this. Nonetheless, will try to have a look at it next week.

[mohdibntarek]

any update here?

[ToucheSir]

Had another look today and surprisingly it seems (k, v) iteration can be handled with just an adjoint (#1285). Iterating values is not yet addressed because that takes an entirely different code path.

[mohdibntarek]

Awesome!

[renatobellotti]

I have a similar small code example (*):

using Zygote

my_dict = Dict(
    "name1" => 1.,
    "name2" => 2.,
)

function my_func(x)
    result = 0
    for (xᵢ, (name, value)) in zip(x, my_dict)
        result += value * xᵢ
    end
    return result
end

my_func([3., 4.])
Zygote.gradient(my_func, [3., 4.])

Error message:

MethodError: no method matching size(::Dict{String, Float64})
Closest candidates are:
  size(::Union{LinearAlgebra.QR, LinearAlgebra.QRCompactWY, LinearAlgebra.QRPivoted}) at <path_edited>/julia-1.7.3/share/julia/stdlib/v1.7/LinearAlgebra/src/qr.jl:567
  size(::Union{LinearAlgebra.QR, LinearAlgebra.QRCompactWY, LinearAlgebra.QRPivoted}, ::Integer) at <path_edited>/julia-1.7.3/share/julia/stdlib/v1.7/LinearAlgebra/src/qr.jl:566
  size(::Union{LinearAlgebra.Cholesky, LinearAlgebra.CholeskyPivoted}) at <path_edited>/julia-1.7.3/share/julia/stdlib/v1.7/LinearAlgebra/src/cholesky.jl:494
  ...

Stacktrace:
  [1] axes
    @ ./abstractarray.jl:95 [inlined]
  [2] _tryaxes(x::Dict{String, Float64})
    @ Zygote ~/.julia/packages/Zygote/D7j8v/src/lib/array.jl:180
  [3] map
    @ ./tuple.jl:222 [inlined]
  [4] adjoint
    @ ~/.julia/packages/Zygote/D7j8v/src/lib/array.jl:293 [inlined]
  [5] _pullback
    @ ~/.julia/packages/ZygoteRules/AIbCs/src/adjoint.jl:65 [inlined]
  [6] _pullback
    @ ./iterators.jl:304 [inlined]
  [7] _pullback(::Zygote.Context, ::typeof(zip), ::Vector{Float64}, ::Dict{String, Float64})
    @ Zygote ~/.julia/packages/Zygote/D7j8v/src/compiler/interface2.jl:0
  [8] _pullback
    @ ./In[49]:8 [inlined]
  [9] _pullback(ctx::Zygote.Context, f::typeof(my_func), args::Vector{Float64})
    @ Zygote ~/.julia/packages/Zygote/D7j8v/src/compiler/interface2.jl:0
 [10] _pullback(f::Function, args::Vector{Float64})
    @ Zygote ~/.julia/packages/Zygote/D7j8v/src/compiler/interface.jl:34
 [11] pullback(f::Function, args::Vector{Float64})
    @ Zygote ~/.julia/packages/Zygote/D7j8v/src/compiler/interface.jl:40
 [12] gradient(f::Function, args::Vector{Float64})
    @ Zygote ~/.julia/packages/Zygote/D7j8v/src/compiler/interface.jl:75
 [13] top-level scope
    @ In[49]:15
 [14] eval
    @ ./boot.jl:373 [inlined]
 [15] include_string(mapexpr::typeof(REPL.softscope), mod::Module, code::String, filename::String)
    @ Base ./loading.jl:1196

Have I understood correctly that iterating over dictionaries is currently not supported?

(*) I know I should not use global variables. This is just a small part of something I'm trying to realise using a closure, but I wanted to simplify the example.

[renatobellotti]

Interestingly, the code works correctly with named tuples (i. e. removing the Dict in the code above). Perhaps this simplifies the bug fix?

[ToucheSir]

Have you tried the linked PR?

[renatobellotti]

No, sorry. I could circumvent the problem by leaving the Dict away.

[CarloLucibello]

closing as it seems #1285 fixed both examples in the OP