Spineax (SParse lINear Solvers in JAX)

This repo integrates existing sparse linear solvers into JAX. I currently feature a single GPU-based linear solver (with plans to implement more):

• cuDSS

For those that need sparsity pattern detection for jax jacobians/hessians I also offer this package.

I built this repo as part of a project to GPU-batch solve many IPOPT optimizations in jaxipm.

cuDSS

I expose most features of cuDSS (as of 0.8.0) to JAX with zero-copy arrays and full FFI jit/vmap integration including custom batching functionality to expose more information than cuDSS currently supports.

This currently supports:

• zero-copies between JAX and cuDSS
• full FFI jit/vmap/grad integration (example)
• all cuDSS datatypes (F32, F64, C64, C128) (example)
• all cuDSS solvers (general, symmetric, symmetric positive defnite, hermitian, hermitian positive definite) (example)
• all cuDSS outputs (example, even in the batched case!)
• Batches of heterogeneous sparsity patterns, and even heterogeneous sizes! (example)

We have also added a new Lineax-based API, which is now the recommended method of interfacing with spineax (example).

Installation

Requirements:

• An NVIDIA GPU of Turing generation (compute capability 7.5) or newer
• CUDA 13
• Python 3.12 or newer
• Linux x86-64 only

pip:

pip install spineax

uv:

uv pip install spineax

Using a uv-managed project instead? Just uv add spineax.

Citation

@article{viljoen2026scaling,
  title={Scaling Nonlinear Optimization: Many Problems One GPU},
  author={Viljoen, John and Haffner, Johanna and Tomizuka, Masayoshi and Mehr, Negar},
  journal={arXiv preprint arXiv:2606.26341},
  year={2026}
}