How to implment `logsumexp` function in JuMP?

Specific Domains Optimization (Mathematical)
, ,
6

And as your advice, I tried to convert it to an equivalent exponential cone program manually but it seems not to work and I’m kinda stuck with it

using ParametrisedConvexApproximators
using JuMP
import DiffOpt
import SCS
import ChainRulesCore
import Flux
using Convex
using UnPack


function main()
    model = Model(() -> DiffOpt.diff_optimizer(SCS.Optimizer))
    empty!(model)
    set_silent(model)
    n, m, d = 3, 2, 1
    i_max = 20
    T = 1e-0
    h_array = [64]
    act = Flux.relu
    plse = PLSE(n, m, i_max, T, h_array, act)
    x = rand(n, d)
    @show plse(x, rand(m, d))
    u_star = infer_test(plse)
    return u_star
end

function infer_test(nn::PLSE)
    n, m, d = 3, 2, 1
    model = Model(() -> DiffOpt.diff_optimizer(SCS.Optimizer))
    @variable(model, u[1:m])

    @unpack T = nn
    x = rand(n, d)
    tmp = ParametrisedConvexApproximators.affine_map(nn, x, u)
    # @NLobjective(model, Min, T * log(sum(exp((1/T)*tmp[i]) for i in 1:nn.i_max)))
    # equivalent form of logsumexp
    _tmp = (1/T) * tmp
    @variable(model, v[1:nn.i_max])
    @variable(model, t)
    @constraint(model, sum(v) <= 1)
    @constraint(model, [i=1:nn.i_max], [v[i], 1, _tmp[i, 1] - t] in MOI.ExponentialCone())
    @objective(model, Min, T*t)
    optimize!(model)
    return value.(u)
end

gives

julia> main()
plse(x, rand(m, d)) = [3.00289608300575;;]
ERROR: Constraints of type MathOptInterface.VectorAffineFunction{Float64}-in-MathOptInterface.ExponentialCone are not supported by the solver.

If you expected the solver to support your problem, you may have an error in your formulation. Otherwise, consider using a different solver.

The list of available solvers, along with the problem types they support, is available at https://jump.dev/JuMP.jl/stable/installation/#Supported-solvers.
Stacktrace:
  [1] error(s::String)
    @ Base ./error.jl:33
  [2] _moi_add_constraint(model::MathOptInterface.Utilities.CachingOptimizer{MathOptInterface.Bridges.LazyBridgeOptimizer{DiffOpt.Optimizer{MathOptInterface.Bridges.LazyBridgeOptimizer{MathOptInterface.Utilities.CachingOptimizer{SCS.Optimizer, MathOptInterface.Utilities.UniversalFallback{MathOptInterface.Utilities.Model{Float64}}}}}}, MathOptInterface.Utilities.UniversalFallback{MathOptInterface.Utilities.Model{Float64}}}, f::MathOptInterface.VectorAffineFunction{Float64}, s::MathOptInterface.ExponentialCone)
    @ JuMP ~/.julia/packages/JuMP/2IF9U/src/constraints.jl:534
  [3] add_constraint(model::Model, con::VectorConstraint{AffExpr, MathOptInterface.ExponentialCone, VectorShape}, name::String)
    @ JuMP ~/.julia/packages/JuMP/2IF9U/src/constraints.jl:560
  [4] macro expansion
    @ ~/.julia/packages/JuMP/2IF9U/src/macros.jl:814 [inlined]
  [5] (::var"#228#232"{VariableRef, Vector{VariableRef}, Matrix{AffExpr}, Model})(i::Int64)
    @ Main ~/.julia/dev/ParametrisedConvexApproximators/test/tmp.jl:27
  [6] #37
    @ ~/.julia/packages/JuMP/2IF9U/src/Containers/container.jl:72 [inlined]
  [7] iterate
    @ ./generator.jl:47 [inlined]
  [8] collect(itr::Base.Generator{JuMP.Containers.VectorizedProductIterator{Tuple{Base.OneTo{Int64}}}, JuMP.Containers.var"#37#38"{var"#228#232"{VariableRef, Vector{VariableRef}, Matrix{AffExpr}, Model}}})
    @ Base ./array.jl:724
  [9] map(f::Function, A::JuMP.Containers.VectorizedProductIterator{Tuple{Base.OneTo{Int64}}})
    @ Base ./abstractarray.jl:2896
 [10] container
    @ ~/.julia/packages/JuMP/2IF9U/src/Containers/container.jl:72 [inlined]
 [11] container
    @ ~/.julia/packages/JuMP/2IF9U/src/Containers/container.jl:66 [inlined]
 [12] infer_test(nn::PLSE)
    @ Main ~/.julia/dev/ParametrisedConvexApproximators/test/tmp.jl:41
 [13] main()
    @ Main ~/.julia/dev/ParametrisedConvexApproximators/test/tmp.jl:23
 [14] top-level scope
    @ REPL[10]:1
 [15] top-level scope
    @ ~/.julia/packages/CUDA/sCev8/src/initialization.jl:52