# How to implment \`logsumexp\` function in JuMP?

**URL:** <https://discourse.julialang.org/t/how-to-implment-logsumexp-function-in-jump/84376>\
**Category:** Optimization (Mathematical)\
**Tags:** jump, differentiation, autodiff\
**Created:** [July 18, 2022, 4:35am UTC](https://discourse.julialang.org/t/how-to-implment-logsumexp-function-in-jump/84376 "2022-07-18T04:35:27Z")\
**Posts on this page:** 8\
**Page:** 1

<div class="post-metadata">

**Author:** ![iHany](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/ihany/32/18151_2.png) [@iHany](https://discourse.julialang.org/u/iHany)\
**Post date:** [July 18, 2022, 4:35am UTC](https://discourse.julialang.org/t/how-to-implment-logsumexp-function-in-jump/84376/1 "2022-07-18T04:35:27Z")

</div>

Hi,  
as I described my situation [here](https://github.com/jump-dev/DiffOpt.jl/issues/228), I’m trying to use DiffOpt.jl for my custom neural network training.

I realised that if I can use `logsumexp` function to express the objective function, I can make it.

However, I can’t find any `logsumexp` for JuMP.  
For example, if I use `Flux.logsumexp`, it says something like this:

```julia
ERROR: MethodError: no method matching isless(::AffExpr, ::AffExpr)
Closest candidates are:
  isless(::Any, ::Missing) at /Applications/Julia-1.7.app/Contents/Resources/julia/share/julia/base/missing.jl:88
  isless(::Missing, ::Any) at /Applications/Julia-1.7.app/Contents/Resources/julia/share/julia/base/missing.jl:87
Stacktrace:
  [1] max(x::AffExpr, y::AffExpr)
    @ Base operators.jl:492
  [2] mapreduce_impl(f::typeof(identity), op::typeof(max), A::Matrix{AffExpr}, first::Int64, last::Int64)
    @ Base reduce.jl:635
  [3] _mapreduce(f::typeof(identity), op::typeof(max), #unused#::IndexLinear, A::Matrix{AffExpr})
    @ Base reduce.jl:417
  [4] _mapreduce_dim(f::typeof(identity), op::typeof(max), #unused#::Base._InitialValue, A::Matrix{AffExpr}, #unused#::Colon)
    @ Base reducedim.jl:330
  [5] mapreduce(f::typeof(identity), op::typeof(max), A::Matrix{AffExpr}; dims::Colon, init::Base._InitialValue)
    @ Base reducedim.jl:322
  [6] mapreduce(f::typeof(identity), op::typeof(max), A::Matrix{AffExpr})
    @ Base reducedim.jl:322
  [7] _maximum(f::typeof(identity), a::Matrix{AffExpr}, ::Colon; kw::Base.Pairs{Symbol, Union{}, Tuple{}, NamedTuple{(), Tuple{}}})
    @ Base reducedim.jl:894
  [8] _maximum(f::typeof(identity), a::Matrix{AffExpr}, ::Colon)
    @ Base reducedim.jl:894
  [9] _maximum(a::Matrix{AffExpr}, ::Colon; kw::Base.Pairs{Symbol, Union{}, Tuple{}, NamedTuple{(), Tuple{}}})
    @ Base reducedim.jl:893
 [10] _maximum(a::Matrix{AffExpr}, ::Colon)
    @ Base reducedim.jl:893
 [11] maximum(a::Matrix{AffExpr}; dims::Colon, kw::Base.Pairs{Symbol, Union{}, Tuple{}, NamedTuple{(), Tuple{}}})
    @ Base reducedim.jl:889
 [12] (::Base.var"#maximum##kw")(::NamedTuple{(:dims,), Tuple{Colon}}, ::typeof(maximum), a::Matrix{AffExpr})
    @ Base reducedim.jl:889
 [13] logsumexp(x::Matrix{AffExpr}; dims::Colon)
    @ NNlib ~/.julia/packages/NNlib/tvMmZ/src/softmax.jl:142
 [14] logsumexp(x::Matrix{AffExpr})
    @ NNlib ~/.julia/packages/NNlib/tvMmZ/src/softmax.jl:142
 [15] infer_test(nn::PLSE)
    @ Main ~/.julia/dev/ParametrisedConvexApproximators/test/tmp.jl:36

```

Is there any workaround for this?

---

<div class="post-metadata">

**Author:** ![odow](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/odow/32/28685_2.png) [@odow](https://discourse.julialang.org/u/odow)\
**Post date:** [July 18, 2022, 5:02am UTC](https://discourse.julialang.org/t/how-to-implment-logsumexp-function-in-jump/84376/2 "2022-07-18T05:02:07Z")

</div>

You typically can’t use arbitrary nonlinear function in JuMP without registering them as user-defined functions: [Nonlinear Modeling · JuMP](https://jump.dev/JuMP.jl/stable/manual/nlp/#User-defined-Functions).

Do you have a reproducible example of what you’re trying to achieve?

The Mosek modeling cookbook [5 Exponential cone optimization — MOSEK Modeling Cookbook 3.3.0](https://docs.mosek.com/modeling-cookbook/expo.html#log-sum-exp) has a conic formulation of logsumexp:

```julia
N = 10
model = Model()
@variable(model, x[1:N])
@NLobjective(model, Min, log(sum(exp(x[i]) for i in 1:N)))
# becomes
model = Model()
@variable(model, x[1:N])
@variable(model, u[1:N])
@variable(model, t)
@constraint(model, sum(u) <= 1)
@constraint(model, [i=1:N], [x[i] - t, 1, u[i]] in MOI.ExponentialCone())
@objective(model, Min, t)

```

---

<div class="post-metadata">

**Author:** ![gdalle](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/gdalle/32/27854_2.png) [@gdalle](https://discourse.julialang.org/u/gdalle)\
**Post date:** [July 18, 2022, 5:07am UTC](https://discourse.julialang.org/t/how-to-implment-logsumexp-function-in-jump/84376/3 "2022-07-18T05:07:05Z")

</div>

An alternative would be to replace JuMP.jl with Convex.jl, which [supports `logsumexp`](https://jump.dev/Convex.jl/stable/operations/#Exponential-Cone-Representable-Functions) (provided you use the right solvers)

---

<div class="post-metadata">

**Author:** ![iHany](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/ihany/32/18151_2.png) [@iHany](https://discourse.julialang.org/u/iHany)\
**Post date:** [July 18, 2022, 5:20am UTC](https://discourse.julialang.org/t/how-to-implment-logsumexp-function-in-jump/84376/4 "2022-07-18T05:20:57Z")

</div>

I thought DiffOpt.jl does not support Convex.jl.  
Is it possible? Do you have any MWE?

---

<div class="post-metadata">

**Author:** ![iHany](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/ihany/32/18151_2.png) [@iHany](https://discourse.julialang.org/u/iHany)\
**Post date:** [July 18, 2022, 5:27am UTC](https://discourse.julialang.org/t/how-to-implment-logsumexp-function-in-jump/84376/5 "2022-07-18T05:27:20Z")

</div>

Could you use this example then?

```julia
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, 1:d])

    @unpack T = nn
    x = rand(n, d)
    tmp = ParametrisedConvexApproximators.affine_map(nn, x, u)
    # T * Flux.logsumexp((1/T)*tmp)
    @NLobjective(model, Min, T * log(sum(exp((1/T)*tmp[i]) for i in 1:nn.i_max)))
    optimize!(model)
    return value.(u)
end

```

---

<div class="post-metadata">

**Author:** ![iHany](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/ihany/32/18151_2.png) [@iHany](https://discourse.julialang.org/u/iHany)\
**Post date:** [July 18, 2022, 5:38am UTC](https://discourse.julialang.org/t/how-to-implment-logsumexp-function-in-jump/84376/6 "2022-07-18T05:38:32Z")

</div>

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

```julia
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
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

```

---

<div class="post-metadata">

**Author:** ![odow](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/odow/32/28685_2.png) [@odow](https://discourse.julialang.org/u/odow)\
**Post date:** [July 18, 2022, 6:20am UTC](https://discourse.julialang.org/t/how-to-implment-logsumexp-function-in-jump/84376/7 "2022-07-18T06:20:57Z")

</div>

Ah. DiffOpt doesnt support exponential cone yet: [Support exponential cone · Issue #50 · jump-dev/DiffOpt.jl · GitHub](https://github.com/jump-dev/DiffOpt.jl/issues/50)

---

<div class="post-metadata">

**Author:** ![iHany](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/ihany/32/18151_2.png) [@iHany](https://discourse.julialang.org/u/iHany)\
**Post date:** [July 18, 2022, 6:23am UTC](https://discourse.julialang.org/t/how-to-implment-logsumexp-function-in-jump/84376/8 "2022-07-18T06:23:24Z")

</div>

Oh no… what a bad news…  
It’s not a proper time to move to Julia for my research yet ☹  
Thx a lot, @odow !