Sparse array primitives and linalg for MLX

[waleed-sh]

I have just released an early beta of mlx-sparse, a small package for sparse array containers and sparse linear algebra primitives in MLX.

The goal is to provide a lightweight MLX-native sparse array package for Apple Silicon. The public API is Python, while performance-critical operations are implemented as MLX primitives with C++/CPU backends and Metal kernels.

The current beta focuses on 2D COO and CSR sparse arrays. It already includes COO to CSR conversion, CSR to dense conversion, CSR canonicalization, CSR matrix-vector products, CSR matrix-matrix products, batched dense RHS products, CSR sparse-sparse products, transpose/Hermitian transpose, and autodiff through sparse values and dense RHS operands.

You can install it from PyPI:

python -m pip install mlx-sparse

This is still an early beta, so APIs may change and some features are incomplete. I would really appreciate feedback from the MLX community, especially bug reports, small reproducible examples, and suggestions for which sparse operations would be most useful next. Please note that this was a personal project which was decided to package together for everyone else, so there may be some quirks.

Edit: it also supports some linalg operations.

More in the docs here.