MathOptAI.jl: Embed trained machine learning predictors into JuMP models

TL;DR

MathOptAI.jl is a Julia library that enables embedding various trained machine learning models like neural networks, decision trees, and Gaussian Processes into optimization models, facilitating hybrid modeling and efficient computation.

Contribution

It introduces a flexible Julia package that integrates multiple ML models into optimization frameworks, including GPU-accelerated evaluations for PyTorch models.

Findings

Supports a wide range of ML models including neural networks, decision trees, and Gaussian Processes.

Enables GPU acceleration for PyTorch model evaluations within optimization.

Provides open-source implementation compatible with Julia and Python ML libraries.

Abstract

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.