Pullback error when summing over a dictionary

[odow]

Is there any guidance on what can and cannot go in an objective function?

Iterating over a dictionary in the objective function seems to break things:

using Nonconvex
Nonconvex.@load Ipopt
data = Dict{Int, Int}(1 => 1);
model = DictModel();
model = addvar!(model, "pg_1", 0.0, Inf);
function opf_objective(x::OrderedDict)
    cost = 0.0
    for (k, v) in data
        cost += x["pg_1"]
    end
    return cost
end
model = set_objective!(model, opf_objective);

julia> optimize(
           model,
           IpoptAlg(),
           NonconvexCore.getinit(model);
           options = IpoptOptions(print_level = 0),
       )
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"#219#220"{Symbol, Dict{Any, Any}})(#unused#::Nothing)
    @ Zygote ~/.julia/packages/Zygote/cCyLF/src/lib/lib.jl:277
  [2] (::Zygote.var"#1770#back#221"{Zygote.var"#219#220"{Symbol, Dict{Any, Any}}})(Δ::Nothing)
    @ Zygote ~/.julia/packages/ZygoteRules/AIbCs/src/adjoint.jl:67
  [3] Pullback
    @ ./Base.jl:34 [inlined]
  [4] (::typeof(∂(setproperty!)))(Δ::Nothing)
    @ Zygote ~/.julia/packages/Zygote/cCyLF/src/compiler/interface2.jl:0
  [5] Pullback
    @ ./dict.jl:679 [inlined]
  [6] (::typeof(∂(skip_deleted_floor!)))(Δ::Nothing)
    @ Zygote ~/.julia/packages/Zygote/cCyLF/src/compiler/interface2.jl:0
  [7] Pullback
    @ ./dict.jl:686 [inlined]
  [8] (::typeof(∂(iterate)))(Δ::Nothing)
    @ Zygote ~/.julia/packages/Zygote/cCyLF/src/compiler/interface2.jl:0
  [9] Pullback
    @ ./REPL[224]:3 [inlined]
 [10] (::typeof(∂(opf_objective)))(Δ::Float64)
    @ Zygote ~/.julia/packages/Zygote/cCyLF/src/compiler/interface2.jl:0
 [11] #207
    @ ~/.julia/packages/Zygote/cCyLF/src/lib/lib.jl:203 [inlined]
 [12] (::Zygote.var"#1747#back#209"{Zygote.var"#207#208"{Tuple{Tuple{Nothing}}, typeof(∂(opf_objective))}})(Δ::Float64)
    @ Zygote ~/.julia/packages/ZygoteRules/AIbCs/src/adjoint.jl:67
 [13] Pullback
    @ ~/.julia/packages/NonconvexCore/YjDSM/src/functions/functions.jl:170 [inlined]
 [14] (::typeof(∂(#_#8)))(Δ::Float64)
    @ Zygote ~/.julia/packages/Zygote/cCyLF/src/compiler/interface2.jl:0
 [15] (::Zygote.var"#207#208"{Tuple{Tuple{Nothing, Nothing}, Tuple{Nothing}}, typeof(∂(#_#8))})(Δ::Float64)
    @ Zygote ~/.julia/packages/Zygote/cCyLF/src/lib/lib.jl:203
 [16] #1747#back
    @ ~/.julia/packages/ZygoteRules/AIbCs/src/adjoint.jl:67 [inlined]
 [17] Pullback
    @ ~/.julia/packages/NonconvexCore/YjDSM/src/functions/functions.jl:170 [inlined]
 [18] (::typeof(∂(λ)))(Δ::Float64)
    @ Zygote ~/.julia/packages/Zygote/cCyLF/src/compiler/interface2.jl:0
 [19] Pullback
    @ ~/.julia/packages/NonconvexCore/YjDSM/src/models/vec_model.jl:90 [inlined]
 [20] (::typeof(∂(λ)))(Δ::Float64)
    @ Zygote ~/.julia/packages/Zygote/cCyLF/src/compiler/interface2.jl:0
 [21] #207
    @ ~/.julia/packages/Zygote/cCyLF/src/lib/lib.jl:203 [inlined]
 [22] (::Zygote.var"#1747#back#209"{Zygote.var"#207#208"{Tuple{Tuple{Nothing}}, typeof(∂(λ))}})(Δ::Float64)
    @ Zygote ~/.julia/packages/ZygoteRules/AIbCs/src/adjoint.jl:67
 [23] Pullback
    @ ~/.julia/packages/NonconvexCore/YjDSM/src/functions/functions.jl:170 [inlined]
 [24] (::typeof(∂(#_#8)))(Δ::Float64)
    @ Zygote ~/.julia/packages/Zygote/cCyLF/src/compiler/interface2.jl:0
 [25] (::Zygote.var"#207#208"{Tuple{Tuple{Nothing, Nothing}, Tuple{Nothing}}, typeof(∂(#_#8))})(Δ::Float64)
    @ Zygote ~/.julia/packages/Zygote/cCyLF/src/lib/lib.jl:203
 [26] #1747#back
    @ ~/.julia/packages/ZygoteRules/AIbCs/src/adjoint.jl:67 [inlined]
 [27] Pullback
    @ ~/.julia/packages/NonconvexCore/YjDSM/src/functions/functions.jl:170 [inlined]
 [28] (::typeof(∂(λ)))(Δ::Float64)
    @ Zygote ~/.julia/packages/Zygote/cCyLF/src/compiler/interface2.jl:0
 [29] Pullback
    @ ~/.julia/packages/NonconvexCore/YjDSM/src/functions/counting_function.jl:9 [inlined]
 [30] (::typeof(∂(λ)))(Δ::Float64)
    @ Zygote ~/.julia/packages/Zygote/cCyLF/src/compiler/interface2.jl:0
 [31] (::Zygote.var"#55#56"{typeof(∂(λ))})(Δ::Float64)
    @ Zygote ~/.julia/packages/Zygote/cCyLF/src/compiler/interface.jl:41
 [32] gradient(f::Function, args::Vector{Float64})
    @ Zygote ~/.julia/packages/Zygote/cCyLF/src/compiler/interface.jl:76
 [33] (::NonconvexIpopt.var"#eval_grad_f#13"{NonconvexCore.CountingFunction{NonconvexCore.Objective{NonconvexCore.var"#133#140"{DictModel, NonconvexCore.Unflatten{OrderedDict{String, Float64}, NonconvexCore.var"#unflatten_to_Dict#64"{typeof(identity), OrderedDict{String, Float64}}}}, Base.RefValue{Float64}, Set{Symbol}}}})(x::Vector{Float64}, grad_f::Vector{Float64})
    @ NonconvexIpopt ~/.julia/packages/NonconvexIpopt/4F7iF/src/ipopt.jl:135
 [34] _Eval_Grad_F_CB(n::Int32, x_ptr::Ptr{Float64}, #unused#::Int32, grad_f::Ptr{Float64}, user_data::Ptr{Nothing})
    @ Ipopt ~/.julia/packages/Ipopt/gZvYd/src/C_wrapper.jl:49
 [35] IpoptSolve(prob::Ipopt.IpoptProblem)
    @ Ipopt ~/.julia/packages/Ipopt/gZvYd/src/C_wrapper.jl:431
 [36] optimize!(workspace::NonconvexIpopt.IpoptWorkspace{NonconvexCore.VecModel{Vector{Float64}}, Ipopt.IpoptProblem, Vector{Float64}, IpoptOptions{NamedTuple{(:print_level, :hessian_approximation, :jac_c_constant, :jac_d_constant), Tuple{Int64, String, String, String}}}, Base.RefValue{Int64}})
    @ NonconvexIpopt ~/.julia/packages/NonconvexIpopt/4F7iF/src/ipopt.jl:52
 [37] #optimize#131
    @ ~/.julia/packages/NonconvexCore/YjDSM/src/models/vec_model.jl:74 [inlined]
 [38] optimize(::DictModel, ::IpoptAlg, ::OrderedDict{String, Float64}; kwargs::Base.Iterators.Pairs{Symbol, IpoptOptions{NamedTuple{(:print_level, :hessian_approximation, :jac_c_constant, :jac_d_constant), Tuple{Int64, String, String, String}}}, Tuple{Symbol}, NamedTuple{(:options,), Tuple{IpoptOptions{NamedTuple{(:print_level, :hessian_approximation, :jac_c_constant, :jac_d_constant), Tuple{Int64, String, String, String}}}}}})
    @ NonconvexCore ~/.julia/packages/NonconvexCore/YjDSM/src/common.jl:233
 [39] top-level scope
    @ REPL[226]:1

cc @ccoffrin

[mohdibntarek]

This looks like a Zygote bug.

f, x = NonconvexCore.getobjective(model), NonconvexCore.getinit(model)
f(x) # works
Zygote.gradient(f, x) # errors

[mohdibntarek]

Nonconvex uses Zygote. Zygote can be surprisingly brittle sometimes.

[ufechner7]

Is there a related Zygote bug report?

[mohdibntarek]

The issue seems to be FluxML/Zygote.jl#1065.

[mohdibntarek]

This is now fixed in the latest release of Zygote.