Sparse arrays/matrices

[shoyer]

Are these in scope for JAX?

I guess the first challenge would be getting sparse arrays into XLA.

[shoyer]

Over in #1032, @Jamesthegiantpeach asked:

Any chance something like this library has been or can be used with jax?

lax_reference is an attempt to implement JAX's core logic on top of NumPy instead of XLA. In principle, I think that means that the JAX could probably work on pydata/sparse via NumPy duck typing (sparse implements NumPy's __array_function__), but on the other hand I'm also pretty sure nobody has tried it and I'm sure this code is quite stale (it does not even appear to have tests).

So one conceivable path for adding pydata/sparse integration (or integration with any other library that implements the NumPy array interface, e.g., dask.array) would be as follows:

1. Add a code transformation for switching a JAX function from lax -> lax_reference inside a function.
2. Add support for missing operations in lax_reference.
3. Tests. Lots of tests.

[jakevdp]

We have some experimental sparse support in JAX that is being actively developed; see https://jax.readthedocs.io/en/latest/jax.experimental.sparse.html

[jakevdp]

I'll close this; you can follow progress in jax.experimental.sparse