MPAX: Mathematical Programming in JAX

MPAX is differentiable via unrolled differentiation. This allows you to integrate mathematical programming directly into the backpropagation process of neural networks. You can use jax.custom_vjp or jax.custom_jvp to define custom forward and backward passes for loss functions that depend on optimization results.

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)

MPAX supports Linear Programming (LP) and Quadratic Programming (QP). You can define problems using either sparse or dense matrix formats for constraints. By default, create_lp and create_qp use sparse matrix formats (use_sparse_matrix=True).

To solve a problem, create the problem object using create_lp or create_qp, instantiate a solver (raPDHG or r2HPDHG), and call .optimize(problem).

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(qp)

MPAX solvers accept several parameters to control execution behavior, convergence, and precision.

General Options

Termination Parameters

Precision and Determinism

Enable Double Precision (64-bit): By default, MPAX uses single-precision (32-bit). To enable 64-bit, run this at the start of your script:

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

Enable Determinism: To ensure deterministic results on GPUs, set the following environment variable:

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

Warning: Do not enable --xla_gpu_deterministic_ops=true when using batch solving, as it significantly degrades performance.

The RestartParameters NamedTuple controls how the algorithm performs restarts to improve convergence.

Key parameters:

• restart_scheme: An instance of RestartScheme (NO_RESTARTS, FIXED_FREQUENCY, or ADAPTIVE_KKT).
• restart_to_current_metric: An instance of RestartToCurrentMetric (NO_RESTART_TO_CURRENT or KKT_GREEDY) to decide between restarting to the average or the current iterate.
• restart_frequency_if_fixed: Frequency of restarts if using FIXED_FREQUENCY.
• artificial_restart_threshold: Fraction of iterations without a restart to trigger an artificial restart.
• sufficient_reduction_for_restart: Threshold improvement to trigger a restart.
• necessary_reduction_for_restart: Necessary threshold improvement to trigger a restart.

from mpax.utils import RestartParameters, RestartScheme, RestartToCurrentMetric

params = RestartParameters(
    restart_scheme=RestartScheme.ADAPTIVE_KKT,
    restart_to_current_metric=RestartToCurrentMetric.KKT_GREEDY,
    restart_frequency_if_fixed=100,
    artificial_restart_threshold=0.5,
    sufficient_reduction_for_restart=0.1,
    necessary_reduction_for_restart=0.01,
    primal_weight_update_smoothing=0.1
)

The TerminationCriteria NamedTuple defines the tolerances used to determine optimality, primal feasibility, and dual feasibility.

When termination_status == OPTIMAL, the following conditions are met:

1. Duality Gap: | primal_objective - dual_objective | <= eps_abs + eps_rel * ( | primal_objective | + | dual_objective | )
2. Primal Feasibility: norm(primal_residual, p) <= eps_abs + eps_rel * norm(right_hand_side, p)
3. Dual Feasibility: norm(dual_residual, p) <= eps_abs + eps_rel * norm(objective_vector, p)

Attributes:

• eps_abs: Absolute tolerance.
• eps_rel: Relative tolerance.
• eps_primal_infeasible: Tolerance for identifying approximate dual rays (certificate of primal infeasibility).
• eps_dual_infeasible: Tolerance for identifying approximate primal rays (certificate of dual infeasibility).
• time_sec_limit: Time limit for the solver.
• iteration_limit: Iteration limit for the solver.

from mpax.utils import TerminationCriteria

criteria = TerminationCriteria(
    eps_abs=1e-6,
    eps_rel=1e-6,
    eps_primal_infeasible=1e-8,
    eps_dual_infeasible=1e-8,
    time_sec_limit=60.0,
    iteration_limit=1000
)

The r2HPDHG class accepts several configuration parameters to tune solver behavior. Below is a reference of the available attributes:

Use create_qp to construct a boxed quadratic program from raw arrays. The function solves:

$$\max \ \frac{1}{2}x^T Q x + c^T x$$ $$\text{s.t. } Ax = b$$ $$\quad \quad Gx \ge h$$ $$\quad \quad l \le x \le u$$

Parameters:

• Q: Quadratic term matrix (supports jnp.ndarray, BCOO, or BCSR).
• c: Objective vector.
• A: Equality constraint matrix.
• b: Right-hand side of equality constraints.
• G: Inequality constraint matrix.
• h: Right-hand side of inequality constraints.
• l: Lower bounds of variables.
• u: Upper bounds of variables.
• use_sparse_matrix: Boolean. If True (default), uses sparse formats for $Q$ and constraints.

Returns:

• A QuadraticProgrammingProblem object.

from mpax.mp_io import create_qp
import jax.numpy as jnp

# Example setup
Q = jnp.array([[1.0, 0.0], [0.0, 1.0]])
c = jnp.array([1.0, 1.0])
A = jnp.array([[1.0, 1.0]])
b = jnp.array([1.0])
G = jnp.array([[1.0, -1.0]])
h = jnp.array([0.0])
l = jnp.array([0.0, 0.0])
u = jnp.array([1.0, 1.0])

qp_problem = create_qp(Q, c, A, b, G, h, l, u)