MPAX: Mathematical Programming in JAX

Mathematical programming has long served as foundation across numerous fields, such as operations research, economics, and engineering, providing powerful robust tools for optimization and decision-making. Recently, these techniques have also found significant applications in machine learning. Notable examples include datadriven decision making [9, 16], learning with physical constraints [8, 10], learning to rank [7], end-to-end planning and control [2], etc. The efficiency and effectiveness of these machine learning approaches depend largely on the rapid processing of large-scale datasets, facilitated by parallel hardware accelerators such as graphics processing units (GPUs). In contrast, traditional approaches to mathematical programming are not well suited for machine learning tasks. Broadly, there are two major paradigms for integrating mathematical programming with machine learning.