SparseTensor#

class iskra.sparse.SparseTensor(tensor: Tensor, *, layout: Literal['coo', 'csr'] = 'coo', dtype: dtype | None = None, device: device | None = None, requires_grad: bool = False)[SOURCE]#

Sane PyTorch sparse tensor.

A subclass of torch.Tensor with a sparse layout, either torch.sparse_coo or torch.sparse_csr. Unlike the default PyTorch sparse tensors, it offers many quality-of-life utilities which make working with sparse tensors a bareable experience in PyTorch, such as scalar/vector/matrix multiplications which just work and ensure your gradients remain sparse (which is somehow not the default in PyTorch), indexing, slicing, as well as small helpers here and there to patch missing functionality.

Primarily, this class supports COO tensors; CSR tensor experience might be mixed.

Warning

Constructing SparseTensor with COO values will automaticall call coalesce() for you, as this is required for alias() to work.

Note

SparseTensor is not necessary to use the free-standing functions in this module. You should be able to use most free-standing functions with a normal tensor created via torch.sparse_coo_tensor. You should be careful with things like multiplications in that case (@ will create dense gradients even if the operands are both sparse matrices).

Important

You almost always want to use sp.coo_tensor() or sp.csr_tensor() to construct a SparseTensor object.

Builds upon and extends existing PyTorch sparse tensors, and offers a (somewhat) sane interface and defaults.

Implementation Details

Alright, I’ve spent a bunch of time digging through PyTorch’s subclassing mess, so I am sharing what I learned here, hoping it is useful in the future. So, our goal is to create a subclass of torch.Tensor which has sparse storage under the hood (i.e., torch.layout == torch.sparse_coo or torch.layout == torch.sparse_csr).

PyTorch’s _make_subclass() is not what we need because it terminates autograd history, meaning, e.g., it won’t propagate gradients to the values of a COO tensor if the values are autograd leafs. In essence, _make_subclass() makes a new leaf tensor with the same data as the input. This is not what we want!

PyTorch’s as_subclass() should be the fix, but fails on sparse tensors. Under the hood, it does a couple of things, see pytorch/pytorch First, it creates a view into the tensor, thus sharing memory but having a new tensor object. Second, it sets __class__ to equal the subclass class. Lastly, it enables __torch_dispatch__ on the subclass via set_python_dispatch(), but only if the subclass has a custom __torch_dispatch__ defined.

Sadly, the PyTorch does not implement aliasing for sparse tensors, but we can simply implement a Python alternative; see alias(). Likewiese, we set tensor.__class__ = SparseTensor manually from Python. The only thing we cannot do from Python is __torch_dispatch__. However, __torch_dispatch__ is only needed if you wish to override low-level tensor behavior, such as behavior during autograd or CUDA kernel calls. In our specific scenario, we are not overriding __torch_dispatch__, so we do not really need to worry about set_python_dispatch(). However, if we were, we would probably need to do somehting akin to this: albanD/subclass_zoo This implementation seems to be more complex, so we will only migrate when this is absolutely necessary.

Constructs a COO SparseTensor from indices and values.

This method wraps torch.sparse_coo_tensor() and does plumbing to convert it to a SparseTensor. The argument and return descriptions reuse the original torch descriptions. See that function for more details.