MPAX: Mathematical Programming in JAX

Haihao Lu, Zedong Peng, Jinwen Yang

Keywords: First-Order Method, Primal-Dual Hybrid Gradient, Linear Programming, Quadratic Programming

Abstract: 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.

Submission Number: 22