Source-linked AI summary
asdex: Automatic Sparse Differentiation in JAX
Adrian Hill, Guillaume Dalle
TL;DR
Dense JAX derivative routines require an AD pass per input or output dimension, with associated memory and computational costs. asdex combines jaxpr-based sparsity detection and graph coloring for sparse differentiation in JAX. It has been applied to Hessian computation in machine learning interatomic potentials and integrated into splineax.
Problem
Dense JAX derivative routines use one AD pass per input or output dimension, creating memory bottlenecks and computational costs that scale with matrix dimension.
Source: Statement of need · p. 2Method
asdex detects dependencies through abstract interpretation of a jaxpr, colors the resulting graph, and prepares a derivative function that composes with jax.jit and jax.vmap.
Source: Core features · p. 3, Performance and interface · p. 4Results
asdex combines jaxpr-based sparsity detection with distance-2 and star coloring for JAX, where the authors found no general-purpose equivalent.
Source: State of the field · pp. 2–3Takeaways & Limitations
asdex has been used to accelerate Hessian computation in machine learning interatomic potentials and integrated into splineax, a sparsity-focused extension of lineax.
Source: Research impact statement · p. 5
Abstract
from arXiv · showhide
Many tasks in scientific computing and machine learning require the Jacobian or Hessian matrix of a function. Automatic differentiation (AD) computes these derivatives to machine precision, but materializing a dense $m \times n$ Jacobian requires $n$ forward-mode or $m$ reverse-mode AD passes, one per column or row. For a large class of functions, each output depends on only a few inputs, making the derivative matrix sparse. Automatic sparse differentiation (ASD) exploits this structure in four steps: detection of the input-agnostic sparsity pattern, coloring of a graph to group columns or rows that can share an AD pass, compressed differentiation to compute a compressed derivative matrix with one AD pass per color, and finally decompression into the original sparsity pattern. The number of colors, and hence of AD passes, is often independent of the problem dimension: a banded Jacobian with $b$ contiguous bands, for instance, only ever requires $b$ colors, regardless of its size. asdex offers the first standalone ASD toolkit in the popular JAX ecosystem. With asdex.jacobian and asdex.hessian, it provides sparse drop-in replacements for jax.jacobian and jax.hessian.
Statement of need
Sparse derivative matrices occur in scientific computing and machine learning, where dense differentiation can impose memory and computational costs. ASD addresses these costs by computing sparse Jacobians or Hessians.
Statement of need
- Sparse derivative matrices arise in nonlinear systems, second-order optimization, sensitivity analysis, and other applications.The target audience for asdex is scientific machine learning researchers and practitioners who want to use JAX performance, JIT compilation, and accelerator support.Source: Statement of need · pp. 1–2
- Dense JAX derivative routines use one AD pass per input or output dimension, creating memory bottlenecks and dimension-scaled computational costs.ASD can alleviate both hurdles through more efficient storage and a faster way to fill sparse matrices.Source: Statement of need · p. 2
State of the field
High-level-language ASD packages and JAX prototypes existed, but the authors found no general-purpose JAX equivalent combining jaxpr-based detection with distance-2 and star coloring. They built asdex to address that gap.
State of the field
- A general-purpose high-level ASD implementation existed in Julia through packages for differentiation, sparsity detection, and coloring.Those packages served as infrastructure for downstream software such as NonlinearSolve.jl.Source: State of the field · p. 2
- No general-purpose ASD equivalent existed for JAX, although several prototypes offered different detection and coloring features.The listed prototypes included distance-1, distance-2, and star coloring, NaN tracing, and jaxpr symbolic conversion.Source: State of the field · p. 2
- Application-specific JAX libraries also implemented parts of the ASD pipeline, including jaxpr-based detection paired with greedy distance-1 or distance-2 coloring.JAX-AMG and tatva were among the concurrent works the authors became aware of while writing the paper.Source: State of the field · p. 3
- asdex combines jaxpr tracing for sparsity detection with distance-2 and star coloring, the combination the authors found missing from existing packages.The passage identifies these as the “proper” detection paradigm and state-of-the-art coloring techniques.Source: State of the field · p. 3
Software design
asdex separates sparsity detection, coloring, compressed differentiation, and decompression into composable stages. Its preparation/evaluation split and JAX-like interface support repeated sparse derivative computation within JAX workflows.
Core features
- Abstract interpretation over a jaxpr propagates input index sets to detect output dependencies, producing a conservative pattern that may include false positives but never misses a nonzero.The same first-order detector obtains Hessian patterns by applying it to jax.grad(f), avoiding a separate second-order detector.Source: Core features · p. 3
- Greedy distance-2 coloring groups Jacobian rows or columns, while star coloring uses Hessian symmetry to reduce colors; asdex defaults to the mode requiring fewer colors for Jacobians.The groups are encoded in a seed matrix.Source: Core features · p. 3
- Seeded Jacobian-vector or vector-Jacobian products compute the compressed matrix, whose nonzero coefficients are scattered into the original sparse format.The products are parallelized with jax.vmap.Source: Core features · pp. 3–4
Performance and interface
- Detection and coloring are one-time preparation based on input shapes, while compressed differentiation and decompression are repeated during evaluation.In iterative solvers or training loops, preparation is amortized; the prepared derivative function composes with jax.jit and jax.vmap and runs on CPU, GPU, and TPU backends.Source: Performance and interface · p. 4
- The API takes a sample input for preparation and returns a function that computes a sparse matrix, defaulting to JAX’s BCOO format.It also supports PyTrees, multiple arguments, auxiliary outputs, known sparsity patterns, reusable colorings, and output formats including NumPy and SciPy sparse formats.Source: Performance and interface · p. 4
- A chunk_size can limit the number of parallel products through jax.lax.map, bounding memory usage.The package provides continuously published benchmarks for detection, coloring, and evaluation performance.Source: Performance and interface · pp. 4–5
Research impact statement
asdex was inspired by a JAX feature request for sparse Jacobian and Hessian support that had remained open since 2019. It has since been applied to accelerate Hessian computation in machine learning interatomic potentials and integrated into splineax.
Research impact statement
- asdex has been applied to accelerate Hessian computation in machine learning interatomic potentials and integrated into splineax, a sparsity-focused extension of lineax.The package was inspired by JAX issue #10329, a sparse Jacobian and Hessian feature request that had remained open since 2019.Source: Research impact statement · p. 5
AI usage disclosure
Generative AI tools were used to develop asdex and write the manuscript; the authors say they reviewed and verified all AI-generated code and text. The software’s design directly builds on the authors’ earlier peer-reviewed Julia packages, which were written without generative AI.
AI usage disclosure
- Generative AI tools, including Anthropic’s Claude Opus 4.5 to 5.5 and Fable 5 models, were used to develop asdex and write the manuscript.Source: AI usage disclosure · p. 5
- The authors reviewed and verified all AI-generated code and text.Source: AI usage disclosure · p. 5
- asdex’s algorithmic approach directly inherits from the authors’ earlier peer-reviewed Julia packages, which were written without generative AI tools.The packages are DifferentiationInterface.jl, SparseConnectivityTracer.jl (Hill & Dalle, 2025), and SparseMatrixColorings.jl (Montoison et al., 2025).Source: AI usage disclosure · p. 5