MathOptAI.jl: Embed trained machine learning predictors into JuMP models
This provides a tool for researchers and practitioners in optimization and machine learning to combine predictive models with optimization frameworks, though it is incremental as it builds on existing libraries and interfaces.
The authors tackled the problem of integrating trained machine learning predictors into mathematical optimization models by developing MathOptAI.jl, a Julia library that embeds various predictors like neural networks and decision trees into JuMP models, with features such as GPU offloading for PyTorch models.
We present \texttt{MathOptAI.jl}, an open-source Julia library for embedding trained machine learning predictors into a JuMP model. \texttt{MathOptAI.jl} can embed a wide variety of neural networks, decision trees, and Gaussian Processes into a larger mathematical optimization model. In addition to interfacing a range of Julia-based machine learning libraries such as \texttt{Lux.jl} and \texttt{Flux.jl}, \texttt{MathOptAI.jl} uses Julia's Python interface to provide support for PyTorch models. When the PyTorch support is combined with \texttt{MathOptAI.jl}'s gray-box formulation, the function, Jacobian, and Hessian evaluations associated with the PyTorch model are offloaded to the GPU in Python, while the rest of the nonlinear oracles are evaluated on the CPU in Julia. \MathOptAI is available at https://github.com/lanl-ansi/MathOptAI.jl under a BSD-3 license.