Add sparse linalg to Jax?

[alewis]

Hi all,

Per the conversation in https://github.com/google/jax/pull/4025 the Jax devs are interested in supporting at least some of the sparse linalg operations we've spent so much time on. Thus we may have more impact by migrating some of our eigs/eigsh/gmres code into Jax directly.