asdex: Automatic Sparse Differentiation in JAX
Abstract
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 × 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.
Related
- Derivative Computation in PINNs: Automatic Differentiation, Finite Differences and Beyond
- Coordinate-Free Automatic Differentiation
- Sparse Identification for Automatic Large-Scale Screening: A Constraint-Aware Framework with Ultra Fast Decoding Algorithm
- Automatic Differentiation with Constraints in Forward and Inverse Problems
- Hyperspectral Anomaly Detection via Group Sparse Low-Rank Tensor Factorization With Automatic Anomaly Grouping
BibTeX
@misc{hill2026asdex,
title = {{asdex: Automatic Sparse Differentiation in JAX}},
author = {Adrian Hill and Guillaume Dalle},
howpublished = {arXiv},
year = {2026},
month = oct,
doi = {10.48550/arxiv.2610.12336},
eprint = {2610.12336},
archivePrefix = {arXiv},
url = {https://arxiv.org/abs/2610.12336},
note = {CC BY 4.0},
}
RIS
TY - UNPB
TI - asdex: Automatic Sparse Differentiation in JAX
AU - Adrian Hill
AU - Guillaume Dalle
PY - 2026
DA - 2026/10/09
PB - arXiv
DO - 10.48550/arxiv.2610.12336
UR - https://arxiv.org/abs/2610.12336
AB - 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 × 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.
ER -