qpax 0.1.0

Differentiable QP solver in JAX.

Differentiable, batched, single-precision quadratic programming in JAX

Features • Installation • Documentation

This package can be used for solving and differentiating (batched) convex quadratic programs of the following form:

$$ \begin{align*} \underset{x}{\text{minimize}} & \quad \frac{1}{2}x^TQx + q^Tx \ \text{s.t.} & \quad Ax = b,\ & \quad Gx \leq h \end{align*} $$

with decision variables $x \in \mathbb{R}^n$, and data matrices $Q \succeq 0$, $q \in \mathbb{R}^n$, $A \in \mathbb{R}^{m \times n}$, $b \in \mathbb{R}^m$, $G \in \mathbb{R}^{p \times n}$ and $h \in \mathbb{R}^p$.

Features

• Differentiable: Backpropagate through QPs and obtain smooth informative subgradients, even at active inequality constraints.
• Single Precision: Runs in f32, allowing for larger batch sizes and higher throughput.
• Batchable: Solves and differentiates lots of QPs in parallel with shared structure.
• Infeasibility avoidance: Avoids generating infeasible problems by solving an always-feasible "elastic" QP and providing informative gradients to encourage feasibility.

Installation

To install directly from github using pip:

• CPU: pip install qpax
• NVIDIA GPU (cuda 12): pip install "qpax[cuda12]"
• NVIDIA GPU (cuda 13): pip install "qpax[cuda13]"

For further details, check our documentation.

License

This project is licensed under the Apache License 2.0 — see the LICENSE file for details.

Citing

If you use this solver, please cite our work(s):

@misc{arrizabalaga2026adifferentiable,
    title         = {A Differentiable Interior-Point Method in Single Precision},
    author        = {Jon Arrizabalaga, Kevin Tracy, Zachary Manchester},
    year          = {2026},
    eprint        = {XXXX},
    archivePrefix = {arXiv},
    primaryClass  = {math.OC}
}

@misc{tracy2024differentiability,
    title={On the Differentiability of the Primal-Dual Interior-Point Method},
    author={Kevin Tracy and Zachary Manchester},
    year={2024},
    eprint={2406.11749},
    archivePrefix={arXiv},
    primaryClass={math.OC}
}