Torch.stack on torch.Tensor subclass

dots: <p>I’m trying to make a tensor subclass that has an extra instance attribute to store data, and I need to store multiple of these subclass objects in a tensor. When I use <code>torch.stack</code> to store them, the elements of the stacked tensor no longer have the extra attribute that they were assigned.</p>
<p>Do you know how we can store multiple of these subclass objects in a tensor (or in a subclass of a tensor) while preserving the extra attributes such that I can access them as <code>stacked[0].extra</code>?</p>
<p>Thank you for your help!</p>
<pre><code class="lang-auto">class tensorExtra(torch.Tensor):
    
    @staticmethod
    def __new__(cls, extra, *args, **kwargs):
        obj = super()._make_subclass(cls, torch.tensor([1,1]), *args, **kwargs)
        obj.__init__(extra)
        return obj
        
    
    def __init__(self, extra):
        self.extra=extra
        
obj1=tensorExtra("abc")
obj2=tensorExtra("def")
list_of_objs=[obj1,obj2]
stacked=torch.stack(list_of_objs)

print(obj1.extra)
print(stacked[0].extra)
</code></pre>

KFrank: <p>Hi Dots!</p>
<aside class="quote no-group" data-username="dots" data-post="1" data-topic="186237" data-full="true">
<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/d/58956e/48.png" class="avatar"> dots:</div>
<blockquote>
<p>I need to store multiple of these subclass objects in a tensor.</p>
</blockquote>
</aside>
<p>The core problem is that a <code>Tensor</code> is in no sense a collection (e.g., a <code>list</code>)<br>
of <code>Tensor</code>s (nor a collection of other objects).  You can’t store a <code>Tensor</code> (nor<br>
another object) “in” a <code>Tensor</code>.</p>
<aside class="quote no-group">
<blockquote>
<p>When I use <code>torch.stack</code> to store them, the elements of the stacked tensor no longer have the extra attribute that they were assigned.</p>
</blockquote>
</aside>
<p>A regular <code>Tensor</code> object does not have an  <code>extra</code> attribute.  But furthermore,<br>
your <code>tensorExtra</code> object does not have per-element <code>extra</code> attributes – it<br>
only carries a single <code>extra</code> attribute for the entire <code>tensorExtra</code> object as a<br>
whole.</p>
<p>And to reiterate what I said above, neither a <code>Tensor</code> nor a <code>tensorExtra</code> is<br>
a collection, so they can’t contain multiple <code>tensorExtra</code>s, each with its own<br>
<code>extra</code> attribute.</p>
<aside class="quote no-group">
<blockquote>
<p>Do you know how we can store multiple of these subclass objects in a tensor</p>
</blockquote>
</aside>
<p>To belabor the point, you can’t, because a <code>Tensor</code> is not a collection in which<br>
you can store multiple objects.  (I doubt it fits the use case you have in mind,<br>
but you certainly can store multiple <code>tensorExtra</code>s in, say, a <code>list</code>, and if you<br>
do, their <code>extra</code> attributes will be preserved and be accessible.)</p>
<aside class="quote no-group">
<blockquote>
<pre><code class="lang-auto">class tensorExtra(torch.Tensor):
    
    @staticmethod
    def __new__(cls, extra, *args, **kwargs):
        obj = super()._make_subclass(cls, torch.tensor([1,1]), *args, **kwargs)
        obj.__init__(extra)
        return obj
        
    
    def __init__(self, extra):
        self.extra=extra
</code></pre>
</blockquote>
</aside>
<p>As an aside, here are some points I do not understand:</p>
<p>As far as I know, <code>torch.stack()</code> is implemented in c++ (with a python wrapper).<br>
Your <code>tensorExtra</code> is a python class about which c++ knows nothing directly.</p>
<p>However, <code>torch.stack()</code> somehow knows – presumably through the python<br>
wrapper machinery – to return a <code>tensorExtra</code> when called with a <code>tensorExtra</code><br>
(albeit a “broken” <code>tensorExtra</code> that lacks an <code>extra</code> attribute).</p>
<p>Consider:</p>
<pre><code class="lang-plaintext">&gt;&gt;&gt; import torch
&gt;&gt;&gt; torch.__version__
'2.0.1'
&gt;&gt;&gt; class tensorExtra(torch.Tensor):
...
...     @staticmethod
...     def __new__(cls, extra, *args, **kwargs):
...         obj = super()._make_subclass(cls, torch.tensor([1,1]), *args, **kwargs)
...         obj.__init__(extra)
...         return obj
...
...
...     def __init__(self, extra):
...         self.extra=extra
...
&gt;&gt;&gt; obj1=tensorExtra("abc")
&gt;&gt;&gt; obj2=tensorExtra("def")
&gt;&gt;&gt; ten3 = torch.tensor ([3, 3])
&gt;&gt;&gt; type (obj1)
&lt;class '__main__.tensorExtra'&gt;
&gt;&gt;&gt; hasattr (obj1, 'extra')
True
&gt;&gt;&gt; type (obj2)
&lt;class '__main__.tensorExtra'&gt;
&gt;&gt;&gt; type (ten3)
&lt;class 'torch.Tensor'&gt;
&gt;&gt;&gt; type (torch.stack ([obj1, obj2]))
&lt;class '__main__.tensorExtra'&gt;
&gt;&gt;&gt; hasattr (torch.stack ([obj1, obj2]), 'extra')
False
&gt;&gt;&gt; type (torch.stack ([obj1, ten3]))
&lt;class '__main__.tensorExtra'&gt;
&gt;&gt;&gt; type (torch.stack ([ten3, obj2]))
&lt;class '__main__.tensorExtra'&gt;
&gt;&gt;&gt; torch.stack ([ten3, obj2])
tensorExtra([[3, 3],
             [1, 1]])
&gt;&gt;&gt; torch.stack ([ten3, obj2])[0]
tensorExtra([3, 3])
&gt;&gt;&gt; type (torch.stack ([ten3, obj2])[0])
&lt;class '__main__.tensorExtra'&gt;
</code></pre>
<p>So, roughly speaking, <code>torch.stack()</code> is trying to implement a covariant return<br>
type when passed* the <code>Tensor</code> subclass <code>tensorExtra</code>.  I’m not sure such a thing<br>
is sound, because (without being overridden) <code>torch.stack()</code> <em>can’t</em> know how to<br>
return a <code>tensorExtra</code>, as evidenced by the fact that the best it can do is to return<br>
a broken <code>tensorExtra</code> that lacks the <code>extra</code> attribute.</p>
<p>I have no idea how or why <code>torch.stack()</code> returns a <code>tensorExtra</code> when a<br>
<code>tensorExtra</code> is passed in.</p>
<p>*)  Of course, the argument to <code>torch.stack()</code> is not a <code>tensorExtra</code> (nor a plain<br>
<code>Tensor</code>).  Rather, the argument is a (python) <code>list</code>.  Now a python <code>list</code> is not a<br>
parameterized type, but to further complicate question of <code>torch.stack()</code>'s return<br>
type, at a conceptual level, should a hypothetical <code>list_of_tensorExtras</code> be a<br>
subtype of <code>list_of_Tensors</code>, be a supertype thereof, be both, or be neither?</p>
<p>Best.</p>
<p>K. Frank</p>

dots: <p>Thank you so much! This is really helpful.</p>
<aside class="quote no-group" data-username="KFrank" data-post="2" data-topic="186237">
<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/k/ecb155/48.png" class="avatar"> KFrank:</div>
<blockquote>
<p>And to reiterate what I said above, neither a <code>Tensor</code> nor a <code>tensorExtra</code> is<br>
a collection, so they can’t contain multiple <code>tensorExtra</code>s, each with its own<br>
<code>extra</code> attribute.</p>
</blockquote>
</aside>
<p>That makes sense. I was getting thrown off by what you mentioned later, where torch.stack was returning a <code>tensorExtra</code> object and, for each “element” in this, <code>type(element)</code> was returning <code>tensorExtra</code>.</p>
<aside class="quote no-group" data-username="KFrank" data-post="2" data-topic="186237">
<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/k/ecb155/48.png" class="avatar"> KFrank:</div>
<blockquote>
<p>I have no idea how or why <code>torch.stack()</code> returns a <code>tensorExtra</code> when a<br>
<code>tensorExtra</code> is passed in.</p>
</blockquote>
</aside>
<p>I’m not sure either. I’ve tried testing this out and editing the class definition, but it somehow returns a <code>tensorExtra</code> with the “broken” <code>extra</code> attribute, which is strange, since  the <code>__new__</code> function explicitly requires an <code>extra</code> attribute to be defined for a <code>tensorExtra</code> to be instantiated.</p>
<p>Thank you again for your help! I appreciate it.</p>