Support for sparse jacobians

[jessegrabowski]

Description

I really want to be able to call pt.jacobian and get back a sparse matrix. For very large systems of equations, the jacobian is essentially never going to be dense, so this would lead to huge performance gains.

I talked with @aseyboldt about it informally once, and he thought it was doable. There is a jax package for it, so that's nice. I confess I don't fully understand what these packages are doing under the hood. From the sparsejac function in the linked package:

    This function uses reverse-mode automatic differentiation to compute the
    Jacobian. The `fn` must accept a rank-1 array and return a rank-1 array, and
    the Jacobian should be sparse with nonzero elements identified by `sparsity`.
    Sparsity is exploited in order to make the Jacobian computation efficient.

    This is done by identifying "structurally independent" groups of output
    elements, which is isomorphic to a graph coloring problem. This allows
    project to a lower-dimensional output space, so that reverse-mode
    differentiation can be more efficiently applied.

Seems like we should be able to do that too!

[jessegrabowski]

I played around a bit with adapting the code from sparsejac. Gist here, it seems easy enough?

[ricardoV94]

This is something where lazy grad ala #788 could be useful, because we can for instance lift the grad_operation (or in this case the sparse_grad) operation above operations like join, so the user doesn't need to artificially define the "distinct" equations ahead of time. It also allows time for pytensor to prune-away zero branches.

This approach may also be combined with our subtensor lift rewrites, like you may make a dense connectivity mask and lift it through the gradient graph, like it is demonstrated in #1471