MPAX: Mathematical Programming in JAX

Haihao Lu, Zedong Peng, Jinwen Yang
12/12/2024
math.OCcs.LG

Official code: MIT-Lu-Lab/MPAX (163 stars, MIT, last commit 2025-06-02).

Abstract

This paper presents MPAX (Mathematical Programming in JAX), a versatile and efficient toolbox for integrating linear programming (LP) into machine learning workflows. MPAX implemented the state-of-the-art first-order methods, restarted average primal-dual hybrid gradient and reflected restarted Halpern primal-dual hybrid gradient, to solve LPs in JAX. This provides native support for hardware accelerations along with features like batch solving, auto-differentiation, and device parallelism. Extensive numerical experiments demonstrate the advantages of MPAX over existing solvers. The solver is available at https://github.com/MIT-Lu-Lab/MPAX.

View on arXivView PDF

Code Implementations(1)

MIT-Lu-Lab/MPAXOfficial100%
MIT

Mathematical Programming in JAX

16311PythonDec 5, 20241 years agoMIT
differentiable-optimizationfirst-order-methodsmathematical-programmingoptimization

Similar Papers

MPX: Mixed Precision Training for JAX

81% match2025

JaxDecompiler: Redefining Gradient-Informed Software Design

79% match1 repo162024

JaxMARL: Multi-Agent RL Environments and Algorithms in JAX

76% match1 repo8622023

JaxPruner: A concise library for sparsity research

76% match1 repo2362023

CLAX: Fast and Flexible Neural Click Models in JAX

75% match1 repo102025

$\texttt{immrax}$: A Parallelizable and Differentiable Toolbox for Interval Analysis and Mixed Monotone Reachability in JAX

75% match1 repo212024

Cite this paper

@article{lu2024mpax,
  title  = {MPAX: Mathematical Programming in JAX},
  author = {Haihao Lu and Zedong Peng and Jinwen Yang},
  year   = {2024},
  eprint = {2412.09734},
  archivePrefix = {arXiv},
  url    = {http://arxiv.org/abs/2412.09734v2}
}

Add a badge to your README

Shows whether official code is linked for this paper.

Code of Paper badge
[![Code of Paper](https://codeofpaper.com/badge/2412.09734.svg)](https://codeofpaper.com/paper/2412.09734)

Discussion