Feature Request: Expose CSC/CSR Compressed Formats and Sparse-Sparse Matrix Multiplication (SpGEMM)

[arubio2]

Describe the feature request
I am working on the GPUmatrix package, which aims to provide drop-in GPU acceleration for R users working with massive sparse matrices (e.g., 20k genes $\times$ 800k spots in spatial transcriptomics).

While the current COO sparse format (torch_sparse_coo_tensor) is great, we are hitting severe performance bottlenecks during matrix algebra. I would like to request exposing/enhancing the bindings for PyTorch's native compressed sparse formats and sparse-to-sparse operations.

Specifically, we need:

1. Compressed Sparse Formats: Full support and documentation for creating and converting to torch.sparse_csc_tensor and torch.sparse_csr_tensor.
2. Sparse $\times$ Sparse Matrix Multiplication (SpGEMM): Enabling torch_mm (or torch_sparse_mm) to accept two sparse matrices. (I believe libtorch currently supports CSR @ CSR and COO @ COO for the forward pass).

Context & Use Case
In iterative algorithms like Non-Negative Matrix Factorization (NMF), we frequently need to multiply large sparse matrices together or perform rapid column/row-wise slicing.

• COO format is too slow for these repeated operations due to uncoalesced memory access.
• Compressed formats (CSC/CSR) natively hook into optimized cuSPARSE routines, providing massive speedups for these exact mathematical workloads.

Expected Behavior
The ability to initialize a compressed sparse tensor in R, and multiply two sparse tensors together, returning a sparse tensor, matching the existing libtorch backend behavior.

Thank you to the mlverse team for all your hard work!