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:

• Lux.Dense
• Lux.Scale
• Lux.relu
• Lux.sigmoid
• Lux.softmax
• Lux.softplus
• Lux.tanh

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, Randomjulia> 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]  ├ -2.7315616607666016 x[1] - moai_Affine[1] = -0.6783187389373779  └ -0.19278962910175323 x[1] - moai_Affine[2] = -0.12212562561035156MathOptAI.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]) = 0Affine(A, b) [input: 2, output: 1]├ variables [1]│ └ moai_Affine[1]└ constraints [1]  └ 0.7000139355659485 moai_ReLU[1] - 1.2195239067077637 moai_ReLU[2] - moai_Affine[1] = 0.04805263504385948

Reduced-space

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

julia> using JuMP, Lux, MathOptAI, Randomjulia> 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.31003424525260925 * max(0, 3.196810245513916 x[1] + 0.3833432197570801) + 0.738827109336853 * max(0, -0.826092004776001 x[1] - 0.8205714225769043)) + 0.5467638373374939julia> 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, Randomjulia> 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.5254679322242737 x[1] - moai_Affine[1] = -0.5781288146972656  └ -1.8065855503082275 x[1] - moai_Affine[2] = 0.1817338466644287MathOptAI.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 [1]  └ 0.232362300157547 moai_ReLU[1] - 0.9425986409187317 moai_ReLU[2] - moai_Affine[1] = 0.11842253059148788