asdex

Automatic Sparse Differentiation in JAX.

asdex (pronounced Aztecs) exploits sparsity structure to efficiently compute sparse Jacobians and Hessians. It implements a custom Jaxpr interpreter that uses abstract interpretation to detect global sparsity patterns from the computation graph, then uses graph coloring to minimize the number of AD passes needed. Refer to our Illustrated Guide to Automatic Sparse Differentiation for more information.

Installation

pip install asdex

Or with uv:

uv add asdex

Quick Example

import asdex
import jax
import numpy as np

def f(x):
    return (x[1:] - x[:-1]) ** 2

x = np.random.randn(1000)

jac_fn = jax.jit(asdex.jacobian(f, x)) # shape and dtype of `x` are used for sparsity pattern detection
J = jac_fn(x)

Instead of 999 VJPs or 1000 JVPs, asdex computes the full sparse Jacobian with just 2 JVPs.

Features

The full ASD pipeline:

• Sparse Jacobians and Hessians: one VJP/JVP/HVP per color, with automatic (or user-defined) mode selection.
• Sparsity detection: finds global sparsity patterns valid for all inputs.
• Graph coloring: row, column, and symmetric coloring minimize AD passes.
• Correctness verification against vanilla JAX.

You already know your sparsity pattern?

• Manually provide sparsity patterns: supply a known pattern from dense, COO, or BCOO formats.
• Precompute, save & load: reuse a colored pattern across inputs, or persist it by saving and loading.

An interface mirroring JAX:

• Multiple inputs and outputs: supports multi-argument functions via argnums, as well as multiple return values.
• PyTree inputs and outputs: sparse differentiation through arbitrary nested PyTrees.
• Auxiliary outputs: supports has_aux=True for functions returning (output, aux).
• Value and derivative: value_and_jacobian / value_and_hessian return the primal value f(x) without a redundant forward pass.

And more:

• Multiple output formats: decompression to BCOO, dense JAX arrays, NumPy, and SciPy (COO/CSR/CSC) arrays.
• Bounded memory: chunk_size caps parallel AD passes for large color counts.
• Visualizations: spy plots and braille pattern previews.

Next Steps

• Getting Started — step-by-step tutorial
• How-To Guides — task-oriented recipes
• Explanation — how and why it works
• API Reference — full API documentation
• Contributing — guidelines for collaborating on asdex
• AI Policy — guidelines for LLM contributions

Acknowledgements

Adrian Hill gratefully acknowledges funding from the German Federal Ministry of Education and Research under the grant BIFOLD26B.

This package is built with Claude Code, based on previous work by Adrian Hill, Guillaume Dalle, and Alexis Montoison in the Julia programming language:

• An Illustrated Guide to Automatic Sparse Differentiation, Hill, Dalle, Montoison (2025)
• Sparser, Better, Faster, Stronger: Efficient Automatic Differentiation for Sparse Jacobians and Hessians, Hill & Dalle (2025)
• Revisiting Sparse Matrix Coloring and Bicoloring, Montoison, Dalle, Gebremedhin (2025)
• SparseConnectivityTracer.jl, Hill, Dalle
• SparseMatrixColorings.jl, Dalle, Montoison
• sparsediffax, Dalle

which in turn stands on the shoulders of giants — notably Andreas Griewank, Andrea Walther, and Assefaw Gebremedhin.

Citation

If you use asdex in your research, please cite:

@software{asdex2026,
  author = {Hill, Adrian},
  title = {asdex: Automatic Sparse Differentiation in JAX},
  url = {https://github.com/adrhill/asdex},
  doi = {10.5281/zenodo.18788242}
}