asdex 0.1.0

Automatic Sparse Differentiation in JAX

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 sparsity patterns from the computation graph, then uses graph coloring to minimize the number of AD passes needed.

[!WARNING] asdex is in early development. The API may change without notice. Use at your own risk.

Installation

pip install git+https://github.com/adrhill/asdex.git

Or with uv:

uv add git+https://github.com/adrhill/asdex.git

Example

import numpy as np
from asdex import jacobian

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

jac_fn = jacobian(f, input_shape=50)
# ColoredPattern(49×50, nnz=98, sparsity=96.0%, JVP, 2 colors)
#   2 JVPs (instead of 49 VJPs or 50 JVPs)
# ⎡⠙⢦⡀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⎤   ⎡⣿⎤
# ⎢⠀⠀⠙⢦⡀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⎥   ⎢⣿⎥
# ⎢⠀⠀⠀⠀⠙⢦⡀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⎥   ⎢⣿⎥
# ⎢⠀⠀⠀⠀⠀⠀⠙⢦⡀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⎥   ⎢⣿⎥
# ⎢⠀⠀⠀⠀⠀⠀⠀⠀⠙⢦⡀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⎥   ⎢⣿⎥
# ⎢⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠙⢦⡀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⎥   ⎢⣿⎥
# ⎢⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠙⢦⡀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⎥ → ⎢⣿⎥
# ⎢⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠙⢦⡀⠀⠀⠀⠀⠀⠀⠀⠀⎥   ⎢⣿⎥
# ⎢⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠙⢦⡀⠀⠀⠀⠀⠀⠀⎥   ⎢⣿⎥
# ⎢⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠙⢦⡀⠀⠀⠀⠀⎥   ⎢⣿⎥
# ⎢⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠙⢦⡀⠀⠀⎥   ⎢⣿⎥
# ⎢⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠙⢦⡀⎥   ⎢⣿⎥
# ⎣⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠉⎦   ⎣⠉⎦

for x in inputs:
    J = jac_fn(x)

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

Documentation

• Getting Started — step-by-step tutorial
• How-To Guides — task-oriented recipes
• Explanation — how and why it works
• API Reference — full API documentation

Acknowledgements

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, A. Hill, G. Dalle, A. Montoison (2025)
• Sparser, Better, Faster, Stronger: Efficient Automatic Differentiation for Sparse Jacobians and Hessians, A. Hill & G. Dalle (2025)
• Revisiting Sparse Matrix Coloring and Bicoloring, A. Montoison, G. Dalle, A. Gebremedhin (2025)
• SparseConnectivityTracer.jl, A. Hill, G. Dalle
• SparseMatrixColorings.jl, G. Dalle, A. Montoison
• sparsediffax, G. Dalle

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