Correct way to extend pytorch sparse Tensor class - C++/CUDA

Hi,
I’m writing an implementation for a new sparse matrix format that i wish to integrate in the pytorch tensor class. What i’m trying to implement is a functionality that works similarly to the one already present for the supported sparse formats, like A.to_sparse_csr(), where A is a tensor. The code is written in C++/CUDA, and although i was able to find something online, i’m struggling a little bit to understand what would be the best way to extend this functionality to python as well. I’m really in the dark at the moment, so i’ll welcome any kind of advice!

Hey!

I would actually suggest looking at a torch_dispatch subclass of Tensor Extending PyTorch — PyTorch main documentation similar to how this “complex Tensor” was done in subclass_zoo/complex_tensor.py at main · albanD/subclass_zoo · GitHub.

These extension points will be significantly simpler than the in-core sparse formats (they didn’t exist when the sparse formats were added). And should allow you to get things running much faster.

For custom kernels, you can see PyTorch Custom Operators — PyTorch Tutorials 2.8.0+cu128 documentation for details (that you can call from your subclass’s torch_dispatch handler for example).