MPAX: Mathematical Programming in JAX

MPAX is a hardware-accelerated, differentiable, batchable, and distributable solver for mathematical programming in JAX, designed to integrate with modern computational and deep learning workflows:

• Hardware accelerated: executes on multiple architectures including CPUs, GPUs and TPUs.
• Differentiable: computes derivatives of solutions with respect to inputs through implicit or unrolled differentiation.
• Batchable: solves multiple problem instances of the same shape simultaneously.
• Distributed: executes in parallel across multiple devices, such as several GPUs.

MPAX's primary motivation is to integrate mathematical programming with deep learning pipelines. To achieve this, MPAX aligns its algorithms and implementations with the requirements of deep learning hardware, ensuring compatibility with GPUs and TPUs. By being differentiable, MPAX can integrate directly into the backpropagation process of neural network training. Its batchability and distributability further enable scalable deployment in large-scale applications.

Currently, MPAX supports linear programming (LP) and quadratic programming (QP), the foundational problems in mathematical programming. Future releases will expand support to include other problem classes of mathematical programming.

Installation

You can install the latest released version of MPAX from PyPI via:

pip install mpax

or you can install the latest development version from GitHub:

pip install git+https://github.com/MIT-Lu-Lab/mpax.git

Quickstart

Currently, MPAX focuses on solving linear programming (LP) and quadratic programming (QP) problems of the following form:

\begin{equation}
\tag{LP}
\begin{aligned}
\min_{l \leq x \leq u}\ & c^\top x \\
\text{s.t.}\ & A x = b \\
& Gx \geq h
\end{aligned}
\end{equation}

\begin{equation}
\tag{QP}
\begin{aligned}
\min_{l \leq x \leq u}\ & \frac{1}{2} x^\top Q x + c^\top x \\
\text{s.t.}\ & A x = b \\
& Gx \geq h
\end{aligned}
\end{equation}

MPAX implements two state-of-the-art first-order methods:

• $\boldsymbol{\mathrm{ra}}$PDHG: restarted average Primal-Dual Hybrid Gradient, supporting both LP (paper) and QP (paper).
• $\boldsymbol{\mathrm{r^2}}$HPDHG: reflected restarted Halpern Primal-Dual Hybrid Gradient, supporting LP only (paper).

Solving a Single LP/QP Problem

MPAX supports both dense and sparse formats for the constraint matrix, controlled by the use_sparse_matrix parameter.

from mpax import create_lp, create_qp, raPDHG, r2HPDHG

# Create LP using sparse matrix format (default)
lp = create_lp(c, A, b, G, h, l, u) # use_sparse_matrix=True by default
# Create LP using dense matrix format
lp = create_lp(c, A, b, G, h, l, u, use_sparse_matrix=False)
solver = r2HPDHG(eps_abs=1e-4, eps_rel=1e-4, verbose=True)
result = solver.optimize(lp)

# Or create QP
qp = create_qp(Q, c, A, b, G, h, l, u)
qp = create_qp(Q, c, A, b, G, h, l, u, use_sparse_matrix=False)
solver = raPDHG(eps_abs=1e-4, eps_rel=1e-4, verbose=True)
result = solver.optimize(lp)

Batch solving

Batch solving allows you to solve multiple LP problems of the same shape simultaneously by using jax.vmap:

import jax.numpy as jnp
from mpax import create_lp, r2HPDHG

def single_optimize(c_vector):
    lp = create_lp(c_vector, A, b, G, h, l, u)
    solver = r2HPDHG(eps_abs=1e-4, eps_rel=1e-4, verbose=True)
    result = solver.optimize(lp)
    obj = jnp.dot(c_vector, result.primal_solution)
    return result.primal_solution, obj

batch_size = 100
batch_c = jnp.tile(c, (batch_size, 1))
batch_optimize = jax.vmap(single_optimize)

result = batch_optimize(batch_c)

Device parallelism

Distribute computations across devices using JAX’s sharding capabilities:

import jax
from mpax import create_lp

# Data sharding
mesh = jax.make_mesh((2,), ('x',))
sharding = jax.sharding.NamedSharding(mesh, P('x',))

A_sharded = jax.device_put(A, sharding)
lp_sharded = create_lp(c, A_sharded, b, G, h, l, u)

solver = r2HPDHG(eps_abs=1e-4, eps_rel=1e-4, verbose=True)
jit_optimize = jax.jit(solver.optimize)
result = jit_optimize(lp_sharded)

Differentiation

An Example of computing the forward and backward passes of the "Smart Predict-then-Optimize+" loss using MPAX and jax.custom_jvp().

import jax
import jax.numpy as jnp

@jax.custom_vjp
def pso_fun(pred_cost, true_cost, true_sol, true_obj):
    sol, obj = batch_optimize(2*pred_cost - true_cost)
    loss = -obj + 2 * jnp.sum(pred_cost * true_sol, axis=1) - true_obj
    loss = jnp.mean(loss)
    return loss, sol

def spo_fwd(pred_cost, true_cost, true_sol, true_obj):
    loss, sol = pso_fun(pred_cost, true_cost, true_sol, true_obj)
    return loss, (sol, true_sol)

def spo_bwd(res, g):
    sol, true_sol = res
    grad = 2 * (true_sol - sol)
    # No gradients needed for true_cost, true_sol, or true_obj
    return grad * g, None, None, None

pso_fun.defvjp(spo_fwd, spo_bwd)

Solver Options

General options

Termination

Precision

By default, MPAX uses single-precision (32-bit). To enable double-precision (64-bit), add the following at the start of your script:

jax.config.update("jax_enable_x64", True)

Determinism

Floating-point computations on GPUs in JAX may produce non-deterministic results. To ensure deterministic results, set:

os.environ["XLA_FLAGS"] = "--xla_gpu_deterministic_ops=true"

Important: If you are using batch solving, do not enable --xla_gpu_deterministic_ops=true, as it can significantly degrade performance.

Citation

If MPAX is useful or relevant to your research, please kindly recognize our contributions by citing our paper:

@article{lu2024mpax,
  title={MPAX: Mathematical Programming in JAX},
  author={Lu, Haihao and Peng, Zedong and Yang, Jinwen},
  journal={arXiv preprint arXiv:2412.09734},
  year={2024}
}