diffqcp
is a JAX library to form the derivative of the solution map to a quadratic cone program (QCP) with respect to the QCP problem data as an abstract linear operator and to compute Jacobian-vector products (JVPs) and vector-Jacobian products (VJPs) with this operator.
The implementation is based on the derivations in our paper (see below) and computes
these products implicitly via projections onto cones and sparse linear system solves.
Our approach therefore differs from libraries that compute JVPs and VJPs by unrolling algorithm iterates.
We directly exploit the underlying structure of QCPs.
Features include:
- Hardware acclerated: JVPs and VJPs can be computed on CPUs, GPUs, and (theoretically) TPUs.
- Support for many canonical classes of convex optimization problems including
- linear programs (LPs),
- quadratic programs (QPs),
- second-order cone programs (SOCPs),
- and semidefinite programs (SDPs).
A quadratic cone program is given by the primal and dual problems
where
diffqcp
currently supports QCPs whose cone is the Cartesian product of the zero cone, the positive orthant, second-order cones, and positive semidefinite cones. Support for exponential and power cones (and their dual cones) is in development.
For more information about these cones, see the appendix of our paper.
diffqcp
is meant to be used as a CVXPYlayers backend --- it is not designed to be a stand-alone
library.
Nonetheless, here is how it use it.
(Note that while we'll specify different CPU and a GPU configurations,
all modules are CPU and GPU compatible--we just recommend the following
as JAX's BCSR
arrays do have CUDA backends for their mv
operations while the BCOO
arrays do not.)
For both of the following problems, we'll use the following objects:
import cvxpy as cvx
problem = cvx.Problem(...)
prob_data, _, _ = problem.get_problem_data(cvx.CLARABEL, solver_opts={'use_quad_obj': True})
scs_cones = cvx.reductions.solvers.conic_solvers.scs_conif.dims_to_solver_dict(prob_data["dims"])
x, y, s = ... # canonicalized solutions to `problem`
If computing JVPs and VJPs on a CPU, we recommend using the equinox.Module
s HostQCP
and QCPStructureCPU
as demonstrated in the following pseudo-example.
from diffqcp import HostQCP, QCPStructureCPU
from jax.experimental.sparse import BCOO
from jaxtyping import Array
P: BCOO = ... # Only the upper triangular part of the QCP matrix P
A: BCOO = ...
q: Array = ...
b: Array = ...
problem_structure = QCPStructureCPU(P, A, scs_cones)
qcp = HostQCP(P, A, q, b, x, y, s, problem_structure)
# Compute JVPs
dP: BCOO ... # Same sparsity pattern as `P`
dA: BCOO = ... # Same sparsity pattern as `A`
db: Array = ...
dq: Array = ...
dx, dy, ds = qcp.jvp(dP, dA, dq, db)
# Compute VJPs
# `dP`, `dA` will be BCOO arrays, `dq`, `db` just Arrays
dP, dA, dq, db = qcp.vjp(f1(x), f2(y), f3(s))
If computing JVPs and VJPs on a GPU, we recommend using the equinox.Module
s QCPStructureGPU
and DeviceQCP
.
from diffqcp import DeviceQCP, QCPStructureGPU
from jax.experimental.sparse import BCSR
from jaxtyping import Array
P: BCSR = ... # The entirety of the QCP matrix P
A: BCSR = ...
q: Array = ...
b: Array = ...
problem_structure = QCPStructureGPU(P, A, scs_cones)
qcp = DeviceQCP(P, A, q, b, x, y, s, problem_structure)
# Compute JVPs
dP: BCSR ... # Same sparsity pattern as `P`
dA: BCSR = ... # Same sparsity pattern as `A`
db: Array = ...
dq: Array = ...
dx, dy, ds = qcp.jvp(dP, dA, dq, db)
# Compute VJPs
# `dP`, `dA` will be BCSR arrays, `dq`, `db` just Arrays
dP, dA, dq, db = qcp.vjp(f1(x), f2(y), f3(s))
@misc{healey2025differentiatingquadraticconeprogram,
title={Differentiating Through a Quadratic Cone Program},
author={Quill Healey and Parth Nobel and Stephen Boyd},
year={2025},
eprint={2508.17522},
archivePrefix={arXiv},
primaryClass={math.OC},
url={https://arxiv.org/abs/2508.17522},
}
diffqcp
is still in development! WIP features and improvements include:
- Support for the exponential cone, the power cone, and their dual cones.
- Batched problem computations.
- Migration of tests from our torch branch.
- Heuristic JVP and VJP computations when the solution map of a QCP is non-differentiable.
Core dependencies (diffqcp
makes essential use of the following libraries)
- Equinox: Neural networks and everything not already in core JAX (via callable
PyTree
s). - Lineax: Linear solvers.
Related
- CVXPYlayers: Construct differentiable convex optimization layers using CVXPY. (WIP:
diffqcp
is being added as a backend for CVXPYlayers.) - CuClarabel: The GPU implemenation of the second-order QCP solver, Clarabel.
- SCS: A first-order QCP solver that has an optional GPU-accelerated backend.
- diffcp: A (Python with C-bindings) library for differentiating through (linear) cone programs.