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> y
1-element Vector{JuMP.VariableRef}: moai_Affine[1]
julia> formulation
Affine(A, b) [input: 1, output: 2] ├ variables [2] │ ├ moai_Affine[1] │ └ moai_Affine[2] └ constraints [2] ├ -3.2024917602539062 x[1] - moai_Affine[1] = 0.14547324180603027 └ -1.76299250125885 x[1] - moai_Affine[2] = -0.688467264175415 MathOptAI.ReLU() ├ variables [2] │ ├ moai_ReLU[1] │ └ moai_ReLU[2] └ constraints [4] ├ moai_ReLU[1] ≥ 0 ├ moai_ReLU[1] - max(0.0, moai_Affine[1]) = 0 ├ moai_ReLU[2] ≥ 0 └ moai_ReLU[2] - max(0.0, moai_Affine[2]) = 0 Affine(A, b) [input: 2, output: 1] ├ variables [1] │ └ moai_Affine[1] └ constraints [2] ├ moai_Affine[1] ≤ 0.2089541256427765 └ -1.1917906999588013 moai_ReLU[1] - 1.1135356426239014 moai_ReLU[2] - moai_Affine[1] = -0.2089541256427765
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> y
1-element Vector{JuMP.NonlinearExpr}: ((+(0.0) + (0.095926433801651 * max(0.0, 2.3050034046173096 x[1] - 0.19422411918640137))) + (-0.8685746788978577 * max(0.0, -2.2948415279388428 x[1] + 0.6971638202667236))) + -0.691771924495697
julia> formulation
ReducedSpace(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> y
1-element Vector{JuMP.VariableRef}: moai_Affine[1]
julia> formulation
Affine(A, b) [input: 1, output: 2] ├ variables [2] │ ├ moai_Affine[1] │ └ moai_Affine[2] └ constraints [2] ├ -1.4951701164245605 x[1] - moai_Affine[1] = 0.2245473861694336 └ -0.7530009746551514 x[1] - moai_Affine[2] = -0.44197702407836914 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 [1] └ 0.8679643869400024 moai_ReLU[1] - 0.23078928887844086 moai_ReLU[2] - moai_Affine[1] = 0.246895894408226