How to inherit sparse tensor

MaKaNu: <p>I try to inherit a sparse tensor.</p>
<p>while the following works:</p>
<pre data-code-wrap="python"><code class="lang-python">class TestTensor(torch.Tensor):
    @staticmethod
    def __new__(cls, data):
        tensor = torch.tensor([data])

        return tensor.as_subclass(cls)
</code></pre>
<p>using instead the <code>sparse_coo_tensor</code> function results in the following situation</p>
<pre data-code-wrap="python"><code class="lang-python">class SparseMasksTensor(torch.Tensor):
    @staticmethod
    def __new__(cls, indices, values, size, requires_grad=False):
        sparse_tensor = torch.sparse_coo_tensor(
            indices, values, size, requires_grad=requires_grad, is_coalesced=True
        )
        return sparse_tensor.as_subclass(cls)

sparse_tensor = SparseMasksTensor([[1,2,3],[0,2,3]], [1,2,3], (4,4))
</code></pre>
<p>Error:</p>
<pre data-code-wrap="python"><code class="lang-python">NotImplementedError: Cannot access storage of SparseTensorImpl
</code></pre>
<p>How can I create inheritance correctly or should I stick to composition in this case?</p>

soulitzer: <blockquote>
<p>How can I create inheritance correctly or should I stick to composition in this case?</p>
</blockquote>
<p>Yeah I’d say that generally composition is preferred when working with PyTorch subclasses.</p>

MaKaNu: <p>That’s sad, since I only extend the Tensor with some methods to make it special…</p>

soulitzer: <blockquote>
<p>That’s sad, since I only extend the Tensor with some methods to make it special…</p>
</blockquote>
<p>Sorry my original reply might’ve been misleading.</p>
<p>To be clear, you’d still be inheriting from a torch.Tensor, except what you want to do is (1) hold the sparse tensor as a field and (2) in your torch dispatch, define that upon every operator call, the operation is passed onto the inner sparse tensor and returned.</p>
<p>Some of the examples here might be helpful <a href="https://github.com/albanD/subclass_zoo" class="inline-onebox" rel="noopener nofollow ugc">GitHub - albanD/subclass_zoo</a></p>

MaKaNu: <aside class="quote no-group" data-username="soulitzer" data-post="4" data-topic="214920">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/letter_avatar_proxy/v4/letter/s/839c29/48.png" class="avatar"> soulitzer:</div>
<blockquote>
<p>(1) hold the sparse tensor as a field and (2) in your torch dispatch, define that upon every operator call, the operation is passed onto the inner sparse tensor and returned.</p>
</blockquote>
</aside>
<p>For future readers: the actual behavior here described can be found in <a href="https://github.com/albanD/subclass_zoo/blob/main/trivial_tensors.py" rel="noopener nofollow ugc">this file</a> in the class <code>TrivialTensorViaComposition</code></p>
<p>EDIT: Following the example I end up again with the issue from above:</p>
<pre data-code-wrap="python"><code class="lang-python">return super().__new__(cls, elem)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^
NotImplementedError: Cannot access storage of SparseTensorImpl
</code></pre>
<p>In this case elem is the sparse implementation. And the error appears in the <code>BaseTensor</code> class, which I left unchanged.</p>

MaKaNu: <p>While the above not worked, the approach from the documentation did kind of:</p>
<p><a href="https://pytorch.org/docs/stable/notes/extending.html#extending-torch-with-a-tensor-wrapper-type" class="onebox" target="_blank" rel="noopener nofollow ugc">https://pytorch.org/docs/stable/notes/extending.html#extending-torch-with-a-tensor-wrapper-type</a></p>
<pre data-code-wrap="python"><code class="lang-python">class MaskSparseTensor(torch.Tensor):
    def __init__(self, data, **kwargs):
        tensor = torch.as_tensor(data, **kwargs)
        self.sparse = tensor.to_sparse_coo()

    def __repr__(self):
        return f"data:\n{self.sparse}"

    @classmethod
    def __torch_function__(cls, func, types, args=(), kwargs=None):
        if kwargs is None:
            kwargs = {}
        args = [getattr(a, "sparse", a) for a in args]
        ret = func(*args, **kwargs)
        return MaskSparseTensor(ret)
</code></pre>
<p>At least I am able now to create instances.</p>
<p>As soon as I try:</p>
<pre data-code-wrap="python"><code class="lang-python">data1 = torch.randn(1000, 3, 4)
data2 = torch.randn(1000, 3, 4)

inst1 = MaskSparseTensor(data1)
inst2 = MaskSparseTensor(data2)

print(inst1 + inst2)
</code></pre>
<p>The well beloved <code>NotImplementedError: Cannot access storage of SparseTensorImpl</code></p>