Lux.jl

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

The upstream documentation is available at https://lux.csail.mit.edu/stable/.

Supported layers

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

Basic example

Use MathOptAI.add_predictor to embed a tuple (containing the Lux.Chain, the parameters, and the state) into a JuMP model:

julia> using JuMP, Lux, MathOptAI, Random
julia> rng = Random.MersenneTwister();
julia> chain = Lux.Chain(Lux.Dense(1 => 2, Lux.relu), Lux.Dense(2 => 1))Chain(
    layer_1 = Dense(1 => 2, relu),                # 4 parameters
    layer_2 = Dense(2 => 1),                      # 3 parameters
)         # Total: 7 parameters,
          #        plus 0 states.
julia> parameters, state = Lux.setup(rng, chain);
julia> predictor = (chain, parameters, state);
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.06411757320165634 x[1] - moai_Affine[1] = 0.9101214408874512
  └ -3.460062026977539 x[1] - moai_Affine[2] = 0.10625481605529785
MathOptAI.ReLU()
├ variables [2]
│ ├ moai_ReLU[1]
│ └ moai_ReLU[2]
└ constraints [4]
  ├ moai_ReLU[1] ≥ 0
  ├ moai_ReLU[1] - max(0, moai_Affine[1]) = 0
  ├ moai_ReLU[2] ≥ 0
  └ moai_ReLU[2] - max(0, moai_Affine[2]) = 0
Affine(A, b) [input: 2, output: 1]
├ variables [1]
│ └ moai_Affine[1]
└ constraints [2]
  ├ moai_Affine[1] ≤ 0.24610388278961182
  └ -0.509725034236908 moai_ReLU[1] - 0.11288736760616302 moai_ReLU[2] - moai_Affine[1] = -0.24610388278961182

Reduced-space

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

julia> using JuMP, Lux, MathOptAI, Random
julia> rng = Random.MersenneTwister();
julia> chain = Lux.Chain(Lux.Dense(1 => 2, Lux.relu), Lux.Dense(2 => 1))Chain(
    layer_1 = Dense(1 => 2, relu),                # 4 parameters
    layer_2 = Dense(2 => 1),                      # 3 parameters
)         # Total: 7 parameters,
          #        plus 0 states.
julia> parameters, state = Lux.setup(rng, chain);
julia> predictor = (chain, parameters, state);
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.49526333808898926 * max(0, 2.547273874282837 x[1] - 0.9145147800445557))) + (1.1176085472106934 * max(0, -0.8526440858840942 x[1] - 0.4929933547973633))) + 0.26910409331321716
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

The Lux extension does not yet support the gray_box keyword argument.

Change how layers are formulated

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

julia> using JuMP, Lux, MathOptAI, Random
julia> rng = Random.MersenneTwister();
julia> chain = Lux.Chain(Lux.Dense(1 => 2, Lux.relu), Lux.Dense(2 => 1))Chain(
    layer_1 = Dense(1 => 2, relu),                # 4 parameters
    layer_2 = Dense(2 => 1),                      # 3 parameters
)         # Total: 7 parameters,
          #        plus 0 states.
julia> parameters, state = Lux.setup(rng, chain);
julia> predictor = (chain, parameters, state);
julia> model = Model();
julia> @variable(model, x[1:1]);
julia> y, formulation = MathOptAI.add_predictor(
           model,
           predictor,
           x;
           config = Dict(Lux.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.41103118658065796 x[1] - moai_Affine[1] = 0.7389950752258301
  └ 2.035665512084961 x[1] - moai_Affine[2] = -0.8400697708129883
MathOptAI.ReLUSOS1()
├ variables [4]
│ ├ moai_ReLU[1]
│ ├ moai_ReLU[2]
│ ├ moai_z[1]
│ └ moai_z[2]
└ constraints [8]
  ├ moai_ReLU[1] ≥ 0
  ├ moai_z[1] ≥ 0
  ├ moai_Affine[1] - moai_ReLU[1] + moai_z[1] = 0
  ├ [moai_ReLU[1], moai_z[1]] ∈ MathOptInterface.SOS1{Float64}([1.0, 2.0])
  ├ moai_ReLU[2] ≥ 0
  ├ moai_z[2] ≥ 0
  ├ moai_Affine[2] - moai_ReLU[2] + moai_z[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.5832075476646423
  └ 0.22390270233154297 moai_ReLU[1] + 1.1555920839309692 moai_ReLU[2] - moai_Affine[1] = -0.5832075476646423