MPAX: Mathematical Programming in JAX

This paper presents MPAX (Mathematical Programming in JAX), a versatile and efficient toolbox for integrating linear programming (LP) and quadratic programming (QP) into machine learning workflows. MPAX implements 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 and QPs in JAX. This provides native support for hardware acceleration, along with features such as batch solving, auto-differentiation, and device parallelism. Extensive numerical experiments demonstrate the advantages of MPAX over existing solvers.