# MPAX: Mathematical Programming in JAX

arXiv: https://arxiv.org/abs/2412.09734 · HTML: https://mlsanity.com/paper/2412.09734/

## Key facts

- Cited by 7 papers in the map (data through 14 Oct 2026).
- Builds on 24 of its 40 references that are in the map.
- First posted on arXiv 12 Dec 2024.

## Builds on (most cited first, 24 of 24)

- [Deep Residual Learning for Image Recognition](https://mlsanity.com/paper/1512.03385/) (2015)
- [PyTorch: An Imperative Style, High-Performance Deep Learning Library](https://mlsanity.com/paper/1912.01703/) (2019)
- [Automatic differentiation in machine learning: a survey](https://mlsanity.com/paper/1502.05767/) (2015)
- [OptNet: Differentiable Optimization as a Layer in Neural Networks](https://mlsanity.com/paper/1703.00443/) (2017)
- [Differentiable Convex Optimization Layers](https://mlsanity.com/paper/1910.12430/) (2019)
- [Smart "Predict, then Optimize"](https://mlsanity.com/paper/1710.08005/) (2017)
- [Differentiation of Blackbox Combinatorial Solvers](https://mlsanity.com/paper/1912.02175/) (2019)
- [Efficient and Modular Implicit Differentiation](https://mlsanity.com/paper/2105.15183/) (2021)
- [Differentiable MPC for End-to-end Planning and Control](https://mlsanity.com/paper/1810.13400/) (2018)
- [OSQP: An Operator Splitting Solver for Quadratic Programs](https://mlsanity.com/paper/1711.08013/) (2017)
- [Practical Large-Scale Linear Programming using Primal-Dual Hybrid Gradient](https://mlsanity.com/paper/2106.04756/) (2021)
- [GPU Acceleration of ADMM for Large-Scale Quadratic Programming](https://mlsanity.com/paper/1912.04263/) (2019)
- [Cardinal Optimizer (COPT) User Guide](https://mlsanity.com/paper/2208.14314/) (2022)
- [Optimization under Uncertainty in the Era of Big Data and Deep Learning: When Machine Learning Meets Mathematical Programming](https://mlsanity.com/paper/1904.01934/) (2019)
- [Faster First-Order Primal-Dual Methods for Linear Programming using Restarts and Sharpness](https://mlsanity.com/paper/2105.12715/) (2021)
- [A Theoretical and Computational Analysis of Full Strong-Branching](https://mlsanity.com/paper/2110.10754/) (2021)
- [A Practical and Optimal First-Order Method for Large-Scale Convex Quadratic Programming](https://mlsanity.com/paper/2311.07710/) (2023)
- [cuPDLP.jl: A GPU Implementation of Restarted Primal-Dual Hybrid Gradient for Linear Programming in Julia](https://mlsanity.com/paper/2311.12180/) (2023)
- [cuPDLP-C: A Strengthened Implementation of cuPDLP for Linear Programming by C language](https://mlsanity.com/paper/2312.14832/) (2023)
- [On the Differentiability of the Primal-Dual Interior-Point Method](https://mlsanity.com/paper/2406.11749/) (2024)
- [Restarted Halpern PDHG for Linear Programming](https://mlsanity.com/paper/2407.16144/) (2024)
- [Restarted Primal-Dual Hybrid Conjugate Gradient Method for Large-Scale Quadratic Programming](https://mlsanity.com/paper/2405.16160/) (2024)
- [HPR-LP: An implementation of an HPR method for solving linear programming](https://mlsanity.com/paper/2408.12179/) (2024)
- [Infeasibility detection with primal-dual hybrid gradient for large-scale linear programming](https://mlsanity.com/paper/2102.04592/) (2021)

## Built on by (most cited first, 7 of 7)

- [Learning Safety Constraints for Large Language Models](https://mlsanity.com/paper/2505.24445/) (2025)
- [OPTIMA: Optimal One-shot Pruning for LLMs via Quadratic Programming Reconstruction](https://mlsanity.com/paper/2512.13886/) (2025)
- [Batched First-Order Methods for Parallel LP Solving in MIP](https://mlsanity.com/paper/2601.21990/) (2026)
- [Contextual Preference Distribution Learning](https://mlsanity.com/paper/2603.17139/) (2026)
- [Parameter Tuning with Generalization Guarantees for GPU-Accelerated Linear Programming](https://mlsanity.com/paper/2606.08638/) (2026)
- [Compression Trinity: Exploring Sparsity, Quantization, and Low-Rank Approximations for LLM Compression](https://mlsanity.com/paper/2608.24070/) (2026)
- [Reinforcement Learning to Accelerate Primal-Dual Hybrid Gradient for Linear Programming](https://mlsanity.com/paper/2610.01546/) (2026)

## Abstract (arXiv)

We present MPAX (Mathematical Programming in JAX), an open-source first-order solver for large-scale linear programming (LP) and convex quadratic programming (QP) built natively in JAX. The primary goal of MPAX is to exploit modern machine learning infrastructure for large-scale mathematical programming, while also providing advanced mathematical programming algorithms that are easy to integrate into machine learning workflows. MPAX implements two PDHG variants, r2HPDHG for LP and rAPDHG for QP, together with diagonal preconditioning, adaptive restarts, adaptive step sizes, primal-weight updates, infeasibility detection, and feasibility polishing. Leveraging JAX's compilation and parallelization ecosystem, MPAX provides across-hardware portability, batched solving, distributed optimization, and automatic differentiation. We evaluate MPAX on CPUs, NVIDIA GPUs, and Google TPUs, observing substantial GPU speedups over CPU baselines and competitive performance relative to GPU-based codebases on standard LP/QP benchmarks. Our numerical experiments further demonstrate MPAX's capabilities in high-throughput batched solving, near-linear multi-GPU scaling for dense LPs, and efficient end-to-end differentiable training. The solver is publicly available at https://github.com/MIT-Lu-Lab/MPAX.