Flux.jl

Flux.jl is a library for machine learning in Julia.

The upstream documentation is available at https://fluxml.ai/Flux.jl/stable/.

Supported layers

MathOptAI supports embedding a Flux model into JuMP if it is a Flux.Chain composed of:

Basic example

Use MathOptAI.add_predictor to embed a Flux.Chain into a JuMP model:

julia> using JuMP, Flux, MathOptAI
julia> predictor = Flux.Chain(Flux.Dense(1 => 2, Flux.relu), Flux.Dense(2 => 1));
julia> model = Model();
julia> @variable(model, x[1:1]);
julia> y, formulation = MathOptAI.add_predictor(model, predictor, x);
julia> y1-element Vector{JuMP.VariableRef}:
 moai_Affine[1]
julia> formulationAffine(A, b) [input: 1, output: 2]
├ variables [2]
│ ├ moai_Affine[1]
│ └ moai_Affine[2]
└ constraints [2]
  ├ 0.5912309288978577 x[1] - moai_Affine[1] = 0
  └ 0.21561956405639648 x[1] - moai_Affine[2] = 0
MathOptAI.ReLU()
├ variables [2]
│ ├ moai_ReLU[1]
│ └ moai_ReLU[2]
└ constraints [4]
  ├ moai_ReLU[1] ≥ 0
  ├ moai_ReLU[2] ≥ 0
  ├ moai_ReLU[1] - max(0.0, moai_Affine[1]) = 0
  └ moai_ReLU[2] - max(0.0, moai_Affine[2]) = 0
Affine(A, b) [input: 2, output: 1]
├ variables [1]
│ └ moai_Affine[1]
└ constraints [1]
  └ -0.9047824144363403 moai_ReLU[1] + 1.3460750579833984 moai_ReLU[2] - moai_Affine[1] = 0

Reduced-space

Use the reduced_space = true keyword to formulate a reduced-space model:

julia> using JuMP, Flux, MathOptAI
julia> predictor = Flux.Chain(Flux.Dense(1 => 2, Flux.relu), Flux.Dense(2 => 1));
julia> model = Model();
julia> @variable(model, x[1:1]);
julia> y, formulation =
           MathOptAI.add_predictor(model, predictor, x; reduced_space = true);
julia> y1-element Vector{JuMP.NonlinearExpr}:
 ((+(0.0) + (-0.6964098811149597 * max(0.0, -1.0810457468032837 x[1]))) + (0.7874495983123779 * max(0.0, -0.10089146345853806 x[1]))) + 0.0
julia> formulationReducedSpace(Affine(A, b) [input: 1, output: 2])
├ variables [0]
└ constraints [0]
ReducedSpace(MathOptAI.ReLU())
├ variables [0]
└ constraints [0]
ReducedSpace(Affine(A, b) [input: 2, output: 1])
├ variables [0]
└ constraints [0]

Gray-box

Use the gray_box = true keyword to embed the network as a nonlinear operator:

julia> using JuMP, Flux, MathOptAI
julia> predictor = Flux.Chain(Flux.Dense(1 => 2, Flux.relu), Flux.Dense(2 => 1));
julia> model = Model();
julia> @variable(model, x[1:1]);
julia> y, formulation =
           MathOptAI.add_predictor(model, predictor, x; gray_box = true);
julia> y1-element Vector{JuMP.VariableRef}:
 moai_GrayBox[1]
julia> formulationGrayBox
├ variables [1]
│ └ moai_GrayBox[1]
└ constraints [1]
  └ op_##610(x[1]) - moai_GrayBox[1] = 0

Change how layers are formulated

Pass a dictionary to the config keyword that maps Flux activation functions to a MathOptAI predictor:

julia> using JuMP, Flux, MathOptAI
julia> predictor = Flux.Chain(Flux.Dense(1 => 2, Flux.relu), Flux.Dense(2 => 1));
julia> model = Model();
julia> @variable(model, x[1:1]);
julia> y, formulation = MathOptAI.add_predictor(
           model,
           predictor,
           x;
           config = Dict(Flux.relu => MathOptAI.ReLUSOS1()),
       );
julia> y1-element Vector{JuMP.VariableRef}:
 moai_Affine[1]
julia> formulationAffine(A, b) [input: 1, output: 2]
├ variables [2]
│ ├ moai_Affine[1]
│ └ moai_Affine[2]
└ constraints [2]
  ├ 0.9946369528770447 x[1] - moai_Affine[1] = 0
  └ 0.41844385862350464 x[1] - moai_Affine[2] = 0
MathOptAI.ReLUSOS1()
├ variables [4]
│ ├ moai_ReLU[1]
│ ├ moai_ReLU[2]
│ ├ moai_z[1]
│ └ moai_z[2]
└ constraints [6]
  ├ moai_ReLU[1] ≥ 0
  ├ moai_ReLU[2] ≥ 0
  ├ moai_Affine[1] - moai_ReLU[1] + moai_z[1] = 0
  ├ moai_Affine[2] - moai_ReLU[2] + moai_z[2] = 0
  ├ [moai_ReLU[1], moai_z[1]] ∈ MathOptInterface.SOS1{Float64}([1.0, 2.0])
  └ [moai_ReLU[2], moai_z[2]] ∈ MathOptInterface.SOS1{Float64}([1.0, 2.0])
Affine(A, b) [input: 2, output: 1]
├ variables [1]
│ └ moai_Affine[1]
└ constraints [2]
  ├ moai_Affine[1] ≥ 0
  └ 0.9569659233093262 moai_ReLU[1] + 1.232018232345581 moai_ReLU[2] - moai_Affine[1] = 0