asdex

Automatic Sparse Differentiation in JAX
Adrian Hill—BIFOLD, TU Berlin Guillaume Dalle—LVMT, ENPC
adrianhill.de
29th EuroAD Workshop, Cambridge, UK
September 29th, 2026
QR code to the slidesSLIDES

Didn't you give this talk already?

Title slide of last year's EuroAD talk
Guillaume Dalle
Alexis Montoison
DifferentiationInterface.jl logo
Klaus-Robert Müller
How about JAX?
Claude logo
Automatic sparse differentiation

Why materialize Jacobians and Hessians?

  • Newton-type solvers
  • Implicit differentiation
  • Laplace approximations
Brusselator Jacobian sparsity pattern
Brusselator PDE Jacobian on an 8×8 grid:
768 of 16384 entries nonzero
JAX basics

What is JAX?

JAX logo
“JAX is a Python library for accelerator-oriented array computation and program transformation, designed for high-performance numerical computing and large-scale machine learning.”

import jax
import jax.numpy as jnp

def f(x):
    y = jnp.sin(x) ** 2
    return jnp.sum(y)

jax.grad(f)(x)
jax.vmap(jax.grad(f))(xs)
jax.jit(jax.vmap(jax.grad(f)))(xs)
JAX docs, The Sharp Bits: JAX has una anima di pura programmazione funzionale

JAX's IR: jaxpr


def f(x):
    y = jnp.sin(x)
    y = y ** 2
    return y
$f(\vx) = \sin(\vx)^2\ \rightarrow$
jax.make_jaxpr(f)(x) 

{ lambda ; a:f32[5]. let
    b:f32[5] = sin a
    c:f32[5] = integer_pow[y=2] b
  in (c,) }
  • explicitly typed
  • functional
  • first-order
  • A-normal form
jax.jvp jax.vmap 

{ lambda ; a:f32[5] b:f32[5]. let
    c:f32[5] = sin a
    d:f32[5] = cos a
    e:f32[5] = mul b d
    f:f32[5] = integer_pow[y=2] c
    g:f32[5] = mul 2.0:f32[] c
    h:f32[5] = mul e g
  in (f, h) }
$\vx,\ \vv\ \rightarrow$
$f(\vx)\ \leftarrow$
$Df(\vx)(\vv)\ \leftarrow$

{ lambda ; a:f32[3,5]. let
    b:f32[3,5] = sin a
    c:f32[3,5] = integer_pow[y=2] b
  in (c,) }

Transformations compose

jax.jvp partial eval jax.linearize transpose jax.vjp seed 1.0 jax.grad
$(\vx, \vv) \mapsto$$(f(\vx),\ Df(\vx)(\vv))$
$\vx \mapsto$$(f(\vx),\ \vv \mapsto Df(\vx)(\vv))$
$\vx \mapsto$$(f(\vx),\ \vw \mapsto Df(\vx)^*(\vw))$
$\vx \mapsto$$Df(\vx)^*(1)$
vmap
jax.jacfwd
$\vx \mapsto \Jf\big|_{\vx}$
vmap
jax.jacrev= jax.jacobian
$\vx \mapsto \Jf\big|_{\vx}$
asdex logo
asdex
github.com/adrhill/asdex

Jacobian sparsity pattern detection demo


import asdex
import jax.numpy as jnp

def f(x):
    y = jnp.sin(x) ** 2
    return y[1:] - y[:-1]

x = jnp.arange(5.0)
print(asdex.jacobian_sparsity(f, x))
      

SparsityPattern(4×5, nnz=8, sparsity=60.0%)
● ● ⋅ ⋅ ⋅
⋅ ● ● ⋅ ⋅
⋅ ⋅ ● ● ⋅
⋅ ⋅ ⋅ ● ●

Hessian sparsity pattern detection


def hessian_sparsity(f, x):
    grad_f = jax.grad(f)
    return jacobian_sparsity(grad_f, x)

For free!

Preparation mechanism

JAXmake_jaxpr(jax.jacfwd(f))(x)

{ lambda ; a:f32[1000]. let
    b:i32[1000,1000] = iota[dimension=0]
    c:i32[1000,1000] = iota[dimension=1]
    d:i32[1000,1000] = add b 0:i32[]
    e:bool[1000,1000] = eq d c
    f:f32[1000,1000] = convert_element_type e
    g:f32[1000,1000] = split f
    h:f32[1000] = sin a
    i:f32[1000] = cos a
    j:f32[1,1000] = broadcast_in_dim i
    k:f32[1000,1000] = mul g j
    l:f32[1000] = integer_pow[y=2] h
    m:f32[1000] = mul 2.0:f32[] h
    n:f32[1,1000] = broadcast_in_dim m
    o:f32[1000,1000] = mul k n
    p:f32[999] = slice[1:1000] l
    q:f32[1000,999] = slice[:,1:1000] o
    r:f32[999] = slice[0:999] l
    s:f32[1000,999] = slice[:,0:999] o
    _:f32[999] = sub p r
    t:f32[1000,999] = sub q s
    u:f32[999,1000] = transpose t
    v:f32[999,1000] = split u
  in (v,) }
asdexmake_jaxpr(asdex.jacobian(f, x))(x)

{ lambda a:f32[2,1000] b:i32[1998,2]
         c:i32[1998,2]; d:f32[1000]. let
    e:f32[1000] = sin d
    f:f32[1000] = cos d
    g:f32[1000] = integer_pow[y=2] e
    h:f32[1000] = mul 2.0:f32[] e
    i:f32[999] = slice[1:1000] g
    j:f32[999] = slice[0:999] g
    _:f32[999] = sub i j
    k:f32[1,1000] = broadcast_in_dim f
    l:f32[2,1000] = mul a k
    m:f32[1,1000] = broadcast_in_dim h
    n:f32[2,1000] = mul l m
    o:f32[2,999] = slice[:,1:1000] n
    p:f32[2,999] = slice[:,0:999] n
    q:f32[2,999] = sub o p
    r:f32[1998] = gather q b
  in (r, c) }
  • Detect + color: slow, once, as in DI.jl [DH26]
  • Seeds $\mS$ are constants
  • One primal pass, shared by all JVPs
  • JVPs batched over 2 colors, not 1000
Just a jaxpr: composes with jit, vmap, grad

Pytree inputs and outputs


def f(x):
    u, p = x["u"], x["p"]
    return {"a": p[0] * u + p[1], "b": jnp.diff(u)}

x = {"u": jnp.array([1.0, 2.0, 3.0]), "p": jnp.array([2.0, 5.0])}
    

jax.jacobian(f)(x)
        

{'a': {'p': Array([[1., 1.],
                   [2., 1.],
                   [3., 1.]], dtype=float32),
       'u': Array([[2., 0., 0.],
                   [0., 2., 0.],
                   [0., 0., 2.]], dtype=float32)},
 'b': {'p': Array([[0., 0.],
                   [0., 0.]], dtype=float32),
       'u': Array([[-1.,  1.,  0.],
                   [ 0., -1.,  1.]], dtype=float32)}}
        

asdex.jacobian(f, x)(x)
        

{'a': {'p': BCOO(float32[3, 2], nse=6),
       'u': BCOO(float32[3, 3], nse=3)},
 'b': {'p': BCOO(float32[2, 2], nse=0),
       'u': BCOO(float32[2, 3], nse=4)}}

Other features

  • Automatic mode: JVPs or VJPs, whichever needs fewer colors
  • Symmetric (star) coloring for Hessians [Geb+07]
  • HVPs as fwd-over-rev, rev-over-fwd or rev-over-rev
  • Control flow: cond, scan, while_loop
  • chunk_size bounds the memory of the batched AD passes
  • argnums, has_aux, value_and_jacobian
  • User-provided sparsity patterns
  • Save and load colorings
  • Output formats: BCOO, dense, NumPy, SciPy
  • Verification against vanilla JAX
  • spy plots of patterns

About robotshumans

  • Tests based on the thorough test suites of SparseConnectivityTracer.jl [HD25] and SparseMatrixColorings.jl [MDG25]
  • Round-trip tests against vanilla JAX: rapid development of sparsity detection
  • Guillaume Dalle reviewed and maintains the coloring code
  • Preprint coming soon
Guillaume Dalle
École nationale des ponts et chaussées, Institut Polytechnique de Paris
LVMT: Laboratoire Ville Mobilité Transport

Applications: truncated sparse AD

Figure 1 of the paper: truncated compressed Hessian evaluation, hop distances in MOF-177, and its Hessian
Title and authors of the truncated automatic sparse differentiation paper
arxiv.org/abs/2609.20510

Bibliography

[Bis+96]C. H. Bischof, P. M. Khademi, A. Bouaricha, and A. Carle, “Efficient Computation of Gradients and Jacobians by Dynamic Exploitation of Sparsity in Automatic Differentiation,” Optimization Methods and Software, vol. 7, no. 1, pp. 1–39, Jan. 1996, doi: 10.1080/10556789608805642. [CM83]T. F. Coleman and J. J. Moré, “Estimation of Sparse Jacobian Matrices and Graph Coloring Problems,” SIAM Journal on Numerical Analysis, vol. 20, no. 1, pp. 187–209, 1983, doi: 10.1137/0720013. [CPR74]A. R. Curtis, M. J. D. Powell, and J. K. Reid, “On the Estimation of Sparse Jacobian Matrices,” IMA Journal of Applied Mathematics, vol. 13, no. 1, pp. 117–119, Feb. 1974, doi: 10.1093/imamat/13.1.117. [DH26]G. Dalle and A. Hill, “A Common Interface for Automatic Differentiation,” Journal of Machine Learning Research, vol. 27, pp. 25:1–25:13, 2026, https://jmlr.org/papers/v27/25-1024.html [DMM90]L. C. W. Dixon, Z. Maany, and M. Mohseninia, “Automatic Differentiation of Large Sparse Systems,” Journal of Economic Dynamics and Control, vol. 14, no. 2, pp. 299–311, May 1990, doi: 10.1016/0165-1889(90)90023-a. [Geb+07]A. H. Gebremedhin, A. Tarafdar, F. Manne, and A. Pothen, “New Acyclic and Star Coloring Algorithms with Application to Computing Hessians,” SIAM Journal on Scientific Computing, vol. 29, no. 3, pp. 1042–1072, 2007, doi: 10.1137/050639879. [GK06]R. Giering and T. Kaminski, “Automatic Sparsity Detection Implemented as a Source-to-Source Transformation,” in Computational Science – ICCS 2006, LNCS, vol. 3994, Springer, 2006, pp. 591–598, doi: 10.1007/11758549_81. [GMP05]A. H. Gebremedhin, F. Manne, and A. Pothen, “What Color Is Your Jacobian? Graph Coloring for Computing Derivatives,” SIAM Review, vol. 47, no. 4, pp. 629–705, 2005, doi: 10.1137/s0036144504444711. [GUG96]U. Geitner, J. Utke, and A. Griewank, “Automatic Computation of Sparse Jacobians by Applying the Method of Newsam and Ramsdell,” in Computational Differentiation: Techniques, Applications, and Tools, Philadelphia, PA: SIAM, 1996, pp. 161–172.

Bibliography

[GW08]A. Griewank and A. Walther, Evaluating Derivatives: Principles and Techniques of Algorithmic Differentiation, 2nd ed. Philadelphia, PA: SIAM, 2008, doi: 10.1137/1.9780898717761. [HD25]A. Hill and G. Dalle, “Sparser, Better, Faster, Stronger: Sparsity Detection for Efficient Automatic Differentiation,” Transactions on Machine Learning Research, Jun. 2025, https://openreview.net/forum?id=GtXSN52nIW [LHC26]M. F. Langer, A. Hill, and M. Ceriotti, “Truncated automatic sparse differentiation for machine learning interatomic potentials,” arXiv:2609.20510, 2026. https://arxiv.org/abs/2609.20510 [MDG25]A. Montoison, G. Dalle, and A. H. Gebremedhin, “Revisiting Sparse Matrix Coloring and Bicoloring,” arXiv:2505.07308, 2025. https://arxiv.org/abs/2505.07308