asdex: Automatic Sparse Differentiation in JAX
Automatic Sparse Differentiation (ASD) efficiently computes sparse Jacobian and Hessian matrices in JAX by exploiting input-dependent sparsity. It achieves this through four steps: detecting input patterns, coloring a graph to group related inputs, performing compressed differentiation with fewer passes, and decompressing the result. This provides a memory-efficient alternative to standard dense differentiation methods.