Regarding #23, I was trying to re-implement this
in Julia.
I found that the use of ComponentArrays may not work properly with Zygote when the optimization parameters are given from another source (here, a neural network) as the following (with v0.5.0-DEV).
I think this is related to jonniedie/ComponentArrays.jl#176.
Therefore, it may be better to clarify the limitation of ComponentArrays for multiple arguments cases (or, just avoid this for now).
Case 1: with ComponentArrays (not working)
Code
using ImplicitDifferentiation
using Flux
using Convex, ECOS
using ComponentArrays
using Random
Random.seed!(2024)
i_max = 10
n = 3
m = 2
model = Chain(
Dense(3, 64, Flux.leakyrelu),
Dense(64, i_max*(m+1)),
)
u_min = -2*ones(m)
u_max = +2*ones(m)
function minimise_logsumexp(θ; T, u_min, u_max, initial_guess, solver=() -> ECOS.Optimizer())
# A = θ[:, 1:end-1]
# B = θ[:, end]
(; A, B) = θ
m = size(A)[2]
u = Convex.Variable(m)
if initial_guess != nothing
u.value = initial_guess
end
obj = T * Convex.logsumexp((1/T)*(A*u + B))
prob = Convex.minimize(obj)
prob.constraints += [u >= u_min]
prob.constraints += [u <= u_max]
solve!(prob, solver(), silent_solver=true, verbose=false)
minimiser = typeof(u.value) <: Number ? [u.value] : u.value[:] # to make it a vector
return minimiser
end
function forward_cstr_optim(θ; kwargs...)
u = minimise_logsumexp(θ; kwargs...)
return u
end
function proj_hypercube(p; u_min, u_max)
return max.(u_min, min.(u_max, p))
end
function conditions_cstr_optim(θ, u; kwargs...)
(; A, B) = θ
# A = θ[:, 1:end-1]
# B = θ[:, end]
∇₂f = A' * Flux.softmax(A*u+B)
η = 0.1
return u .- proj_hypercube(u .- η .* ∇₂f; kwargs...)
end
function implicit_cstr_optim_func(θ; T, u_min, u_max, initial_guess, solver)
tmp = ImplicitFunction(
θ -> forward_cstr_optim(θ; T, u_min, u_max, initial_guess, solver),
(θ, u) -> conditions_cstr_optim(θ, u; u_min, u_max),
)
tmp(θ)
end
x = 100*(2*rand(n) .- 1)
val, grads = Flux.withgradient(model) do _model
AB = reshape(_model(x), i_max, m+1)
A = AB[:, 1:m]
B = AB[:, m+1]
θ = ComponentArray(A=A, B=B)
# θ = reshape(_model(x), i_max, m+1)
u_star = implicit_cstr_optim_func(θ; T=1, u_min=-ones(m), u_max=+ones(m), initial_guess=nothing, solver=() -> ECOS.Optimizer())
sum(u_star)
endResult
ERROR: Mutating arrays is not supported -- called push!(Vector{Any}, ...)
This error occurs when you ask Zygote to differentiate operations that change
the elements of arrays in place (e.g. setting values with x .= ...)
Possible fixes:
- avoid mutating operations (preferred)
- or read the documentation and solutions for this error
https://fluxml.ai/Zygote.jl/latest/limitationsCase 2: without ComponentArrays (working)
Code
using ImplicitDifferentiation
using Flux
using Convex, ECOS
using ComponentArrays
using Random
Random.seed!(2024)
i_max = 10
n = 3
m = 2
model = Chain(
Dense(3, 64, Flux.leakyrelu),
Dense(64, i_max*(m+1)),
)
u_min = -2*ones(m)
u_max = +2*ones(m)
function minimise_logsumexp(θ; T, u_min, u_max, initial_guess, solver=() -> ECOS.Optimizer())
# A = θ[:, 1:end-1]
# B = θ[:, end]
(; A, B) = θ
m = size(A)[2]
u = Convex.Variable(m)
if initial_guess != nothing
u.value = initial_guess
end
obj = T * Convex.logsumexp((1/T)*(A*u + B))
prob = Convex.minimize(obj)
prob.constraints += [u >= u_min]
prob.constraints += [u <= u_max]
solve!(prob, solver(), silent_solver=true, verbose=false)
minimiser = typeof(u.value) <: Number ? [u.value] : u.value[:] # to make it a vector
return minimiser
end
function forward_cstr_optim(θ; kwargs...)
u = minimise_logsumexp(θ; kwargs...)
return u
end
function proj_hypercube(p; u_min, u_max)
return max.(u_min, min.(u_max, p))
end
function conditions_cstr_optim(θ, u; kwargs...)
# (; A, B) = θ
A = θ[:, 1:end-1]
B = θ[:, end]
∇₂f = A' * Flux.softmax(A*u+B)
η = 0.1
return u .- proj_hypercube(u .- η .* ∇₂f; kwargs...)
end
function implicit_cstr_optim_func(θ; T, u_min, u_max, initial_guess, solver)
tmp = ImplicitFunction(
θ -> forward_cstr_optim(θ; T, u_min, u_max, initial_guess, solver),
(θ, u) -> conditions_cstr_optim(θ, u; u_min, u_max),
)
tmp(θ)
end
x = 100*(2*rand(n) .- 1)
val, grads = Flux.withgradient(model) do _model
AB = reshape(_model(x), i_max, m+1)
# A = AB[:, 1:m]
# B = AB[:, m+1]
# θ = ComponentArray(A=A, B=B)
θ = reshape(_model(x), i_max, m+1)
u_star = implicit_cstr_optim_func(θ; T=1, u_min=-ones(m), u_max=+ones(m), initial_guess=nothing, solver=() -> ECOS.Optimizer())
sum(u_star)
endResult
(val = 1.2111089906277515, grad = ((layers = ((weight = Float32[-1.9894879 0.016161848 -1.9257156; -0.009999948 8.12358f-5 -0.009679403; … ; 0.31266153 -0.0025399441 0.30263928; -1.2044753 0.009784702 -1.1658663], bias = Float32[-0.026939616, -0.0001354091, 0.0067554815, -0.012562656, -0.0116912015, -0.0072106766, -0.019949805, -0.012255933, -0.008887393, -0.00019493952 … 8.364284f-6, 0.0029845645, 0.00020423913, -0.00016008505, -0.00011763882, 0.01743382, 0.015867122, 0.013238616, 0.0042337435, -0.016309775], σ = nothing), (weight = Float32[2.7479534f-31 -8.035564f-33 … 6.655762f-31 8.295855f-31; 4.978378f-17 -1.455777f-18 … 1.2058026f-16 1.502933f-16; … ; 5.054912f-10 -1.478157f-11 … 1.2243399f-9 1.5260381f-9; 3.615537f-16 -1.0572551f-17 … 8.757118f-16 1.0915021f-15], bias = Float32[2.8529254f-32, 5.168552f-18, 1.702376f-20, -0.011397834, 0.012743557, 3.4305703f-10, 2.0157113f-23, -0.010961587, 9.4864516f-12, 7.707587f-18 … 1.4720016f-31, 2.3545986f-17, 7.705814f-20, -0.035908565, 0.057037998, 1.6192033f-9, 9.8762665f-23, -0.021129435, 5.2480097f-11, 3.7536507f-17], σ = nothing)),),))
![github-actions[bot] avatar](https://avatars.githubusercontent.com/in/15368?v=4 "github-actions[bot]")
React
Typescript
Django
Laravel
Facebook
Microsoft
Google
Alibaba
D3
Tencent