Part of the reason we implemented a lot of how we handle W ourselves is because back in like PyTorch 0.2 when we first implemented SKI, the sparse tensor support in native PyTorch was pretty slow.

It's still not a guarantee that just implementing W as a sparse tensor in native torch will be faster (or possible -- back then they didn't support batch matmuls), because W has special additional structure, but it could be worth a look.

If it's still better to special case W, I agree we should have SparseLazyTensor that wraps PyTorch native sparse tensors, and perhaps a subclass that deals with W specifically.

Finally, I wonder if this effort should be done on the linear_operator package we're trying to pull out of GPyTorch, because SparseLazyTensors aren't necessarily positive definite. This could be a good motivation to finally finish up a minimal version of that package and make it a dependency for GPyTorch.

Thanks for commenting @jacobrgardner.

I'll try to run some simple tests that shall allow us to compare current implementation with torch sparse case.

For now, I'll push this work here only as I'm not able to see how gpytorch depends on the linera_operator library. Feel free to drop any pointer if I'm missing any infomation.

Linear_operator lives here: https://github.com/cornellius-gp/linear_operator and as you can see it's currently an inconsistently updated clone of gpytorch.lazy but the goal is to have linear_operator be a (more) stand-alone linear algebra package that supports generic linear operators (e.g. lazy tensors) beyond just those useful for GPs.

Ultimately the hope is that the basic infrastructure for this will be implemented in pytorch proper (as it is for scipy / Tensorflow): pytorch/pytorch#28341.