arXiv preprint ·
asdex: Automatic Sparse Differentiation in JAX
- Published
- Source
- arXiv preprint, posted before peer review
- Original
- Read it at arXiv
- Topics
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 \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 $\texttt{asdex.jacobian}$ and $\texttt{asdex.hessian}$, it provides sparse drop-in replacements for $\texttt{jax.jacobian}$ and $\texttt{jax.hessian}$.
The authors’ abstract, as arXiv lists it. arXiv’s listings, abstracts included, are free for anyone to reuse.
Read it on ReadSearch
Each opens after you create an account, and brings you back to this paper.
More in Mathematical Software
Separating Statistical Inference from Experimental Runtime: An Architecture for Sequential Scientific Analysis
Coordinate-Free Automatic Differentiation
GpFlash: Maximizing Speed and Minimizing Memory for Gaussian Process Training and Prediction
Scen-Opt: A Scenario Optimization Toolbox for Data-Driven Convex Programming
Learning to Design Convex Relaxations for Nonconvex Optimization
Get the few papers worth your time in Mathematical Software every morning.
ReadSearch reads the new papers in your field each day and brings you the few worth reading, each with its key finding in a sentence.