Extending torch.Tensor

username: nitaifingerhut
cooked: <p>Hi all,<br>
I’m trying to extend <code>torch.Tensor</code> as follow:</p>
<pre><code class="lang-auto">class MyTensor1(torch.Tensor):
    @staticmethod
    def __new__(cls, data, **kwargs):
        tensor = torch.tensor(data, **kwargs)

        if tensor.shape != (3, 3):
            raise ValueError

        return super().__new__(cls, data, **kwargs)

    @property
    def mid(self):
        return self[1, 1]

    @mid.setter
    def mid(self, val: float):
        self[1, 1] = val * torch.ones_like(self[1, 1])
</code></pre>
<p>and</p>
<pre><code class="lang-auto">class MyTensor2(torch.Tensor):
    @staticmethod
    def __new__(cls, data, **kwargs):
        tensor = torch.tensor(data, **kwargs)

        if tensor.shape != (3, 3):
            raise ValueError

        return super().__new__(cls, data, **kwargs)

    @property
    def mid_low(self):
        return self[1, 2]

    @ mid_low
    def mid_low(self, val: float):
        self[1, 2] = val * torch.ones_like(self[1, 2])
</code></pre>
<p>following that I have two questions/problems:</p>
<ol>
<li>I want to be able to return a <code>torch.Tensor</code> object from accessing properties or any sort of slicing. Currently:</li>
</ol>
<pre><code class="lang-auto">tensor1 = MyTensor1([[1,1,1],[2,2,2],[3,3,3]])
x = tensor1.mid
y = tensor1[:, -1]
</code></pre>
<p>the types of resulting <code>x</code> or <code>y</code> are <code>MyTensor1</code>. Is it possible that it will be <code>torch.Tensor</code> (without explicitly using <code>torch.tensor()</code> here or in the property call)?</p>
<ol start="2">
<li>when I try to apply some operations, it seems like inheritance wasn’t good. for example:</li>
</ol>
<pre><code class="lang-auto">tensor1 = MyTensor1([[1,1,1],[2,2,2],[3,3,3]])
tensor2 = MyTensor2([[1,1,1],[2,2,2],[3,3,3]])
x = tensor1 @ tensor2
</code></pre>
<p>resulting</p>
<pre><code class="lang-auto">TypeError: unsupported operand type(s) for @: 'MyTensor1' and 'MyTensor2'
</code></pre>
<p>even though the operation is completely defined.</p>
<p>Thanks for your help!</p>

username: nitaifingerhut
cooked: <p>Edit:<br>
from docs:<br>
<a href="https://pytorch.org/docs/stable/notes/extending.html" class="onebox" target="_blank" rel="noopener nofollow ugc">https://pytorch.org/docs/stable/notes/extending.html</a></p>
<p>I tried to add</p>
<pre><code class="lang-auto">class MyTensor1(torch.Tensor):
    @staticmethod
    def __new__(cls, data, **kwargs):
        tensor = torch.tensor(data, **kwargs)

        if tensor.shape != (3, 3):
            raise ValueError

        return super().__new__(cls, data, **kwargs)

    @classmethod
    def __torch_function__(cls, func, types, args=(), kwargs=None):
        if kwargs is None:
            kwargs = {}
        return super().__torch_function__(func, types, args, kwargs)

    @property
    def mid(self):
        return self[1, 1]

    @mid.setter
    def mid(self, val: float):
        self[1, 1] = val * torch.ones_like(self[1, 1])
</code></pre>
<p>but it didn’t solve issue <span class="hashtag">#2</span></p>
<p>As for issue <span class="hashtag">#1:</span><br>
I opened a <a href="https://github.com/pytorch/pytorch/issues/69893" rel="noopener nofollow ugc">feature request</a>, but wondering if there is an alternative for now…</p>