Differentiable matrix-free linear algebra, optimization and equation solving

This is a meta-issue for keeping track of progress on implementing differentiable higher-order functions from SciPy in JAX, e.g.,

• scipy.sparse.linalg.gmres and cg: matrix-free linear solves
• scipy.sparse.linalg.eigs and eigsh: matrix-free eigenvalue problems
• scipy.optimize.root: nonlinear equation solving
• scipy.optimize.fixed_point: solving for fixed points
• scipy.integrate.odeint: solving ordinary differential equations
• scipy.optimize.minimize: nonlinear minimization

These higher-order functions are important for implementing sophisticated differentiable programs, both for scientific applications and for machine learning.

Implementations should leverage and build upon JAX’s custom transformation capabilities. For example, scipy.optimize.root should leverage autodiff for calculating the Jacobians or Jacobian-vector products needed for Newton’s method.

In most cases, I think the right way to do this involves two separate steps, which could happens in parallel:

1. Higher order primitives for defining automatic differentiation rules, but not specialized to any particular algorithm, e.g., lax.custom_linear_solve from https://github.com/google/jax/pull/1402.
2. Implementations of particular algorithms for the forward problems, e.g., a conjugate gradient method for linear solves. These could either be implemented from scratch using JAX’s functional control flow (e.g., while_loop) or could leverage existing external implementations on particular backends. Either way they will almost certainly need custom derivative rules, rather than differentiation through the forward algorithm.

There’s lots of work to be done here, so please comment if you’re interested in using or implementing any of these.

[mattjj]

@romanodev this is really awesome work.

[romanodev]

@shoyer thanks for sharing! I think it would be nice to combine your implementation with the dot product between a sparse matrix and a vector #3717. The jit/GPU implementation still can’t beat Scipy and I suspect this is due to the COO representation of the sparse matrix (Scipy uses CSR https://github.com/scipy/scipy/blob/v1.5.1/scipy/sparse/base.py#L532). I will do some testing in this direction first.