Prototype for sparse automatic differentiation in JAX + Julia.
This package is meant as an alternative to sparsejac with the following differences:
- More efficient graph encodings and colorings thanks to the Julia library SparseMatrixColorings.jl
- Optimized symmetry-aware computation of sparse Hessians (although taking the sparse Jacobian of the gradient can also give good results in practice)
See the documentation for details on the API.
Warning
This is a work in progress, it needs more docs and tests. Try at your own risk.