mlx
mlx copied to clipboard
[Experiment] Kernel for sparse matrix-dense matrix multiplication
Proposed changes
Feature: added function metal kernel sparse_matmul_csr for sparse matrix-dense matrix multiplication. The function takes as input a sparse matrix (represented as three arrays representing the row pointers, the column indices, and the values at the nonzero entries) and a dense matrix represented as a dense array with two dimensions. I have implemented the backends for metal and for the cpu but not cuda.
Please let me know if there is any documentation I should update.
Checklist
Put an x in the boxes that apply.
- [x] I have read the CONTRIBUTING document
- [x] I have run
pre-commit run --all-filesto format my code / installed pre-commit prior to committing changes - [x] I have added tests that prove my fix is effective or that my feature works
- [ ] I have updated the necessary documentation (if needed)