Torch.stack on torch.Tensor subclass

[dots]

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 torch.stack to store them, the elements of the stacked tensor no longer have the extra attribute that they were assigned.

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 stacked[0].extra?

Thank you for your help!

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)

[KFrank]

Hi Dots!

@dots

I need to store multiple of these subclass objects in a tensor.

The core problem is that a Tensor is in no sense a collection (e.g., a list)
of Tensors (nor a collection of other objects). You can’t store a Tensor (nor
another object) “in” a Tensor.

When I use torch.stack to store them, the elements of the stacked tensor no longer have the extra attribute that they were assigned.

A regular Tensor object does not have an extra attribute. But furthermore,
your tensorExtra object does not have per-element extra attributes – it
only carries a single extra attribute for the entire tensorExtra object as a
whole.

And to reiterate what I said above, neither a Tensor nor a tensorExtra is
a collection, so they can’t contain multiple tensorExtras, each with its own
extra attribute.

Do you know how we can store multiple of these subclass objects in a tensor

To belabor the point, you can’t, because a Tensor is not a collection in which
you can store multiple objects. (I doubt it fits the use case you have in mind,
but you certainly can store multiple tensorExtras in, say, a list, and if you
do, their extra attributes will be preserved and be accessible.)

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

As an aside, here are some points I do not understand:

As far as I know, torch.stack() is implemented in c++ (with a python wrapper).
Your tensorExtra is a python class about which c++ knows nothing directly.

However, torch.stack() somehow knows – presumably through the python
wrapper machinery – to return a tensorExtra when called with a tensorExtra
(albeit a “broken” tensorExtra that lacks an extra attribute).

Consider:

>>> import torch
>>> torch.__version__
'2.0.1'
>>> 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")
>>> ten3 = torch.tensor ([3, 3])
>>> type (obj1)
<class '__main__.tensorExtra'>
>>> hasattr (obj1, 'extra')
True
>>> type (obj2)
<class '__main__.tensorExtra'>
>>> type (ten3)
<class 'torch.Tensor'>
>>> type (torch.stack ([obj1, obj2]))
<class '__main__.tensorExtra'>
>>> hasattr (torch.stack ([obj1, obj2]), 'extra')
False
>>> type (torch.stack ([obj1, ten3]))
<class '__main__.tensorExtra'>
>>> type (torch.stack ([ten3, obj2]))
<class '__main__.tensorExtra'>
>>> torch.stack ([ten3, obj2])
tensorExtra([[3, 3],
             [1, 1]])
>>> torch.stack ([ten3, obj2])[0]
tensorExtra([3, 3])
>>> type (torch.stack ([ten3, obj2])[0])
<class '__main__.tensorExtra'>

So, roughly speaking, torch.stack() is trying to implement a covariant return
type when passed* the Tensor subclass tensorExtra. I’m not sure such a thing
is sound, because (without being overridden) torch.stack() can’t know how to
return a tensorExtra, as evidenced by the fact that the best it can do is to return
a broken tensorExtra that lacks the extra attribute.

I have no idea how or why torch.stack() returns a tensorExtra when a
tensorExtra is passed in.

*) Of course, the argument to torch.stack() is not a tensorExtra (nor a plain
Tensor). Rather, the argument is a (python) list. Now a python list is not a
parameterized type, but to further complicate question of torch.stack()'s return
type, at a conceptual level, should a hypothetical list_of_tensorExtras be a
subtype of list_of_Tensors, be a supertype thereof, be both, or be neither?

Best.

K. Frank

[dots]

Thank you so much! This is really helpful.

@KFrank

And to reiterate what I said above, neither a Tensor nor a tensorExtra is
a collection, so they can’t contain multiple tensorExtras, each with its own
extra attribute.

That makes sense. I was getting thrown off by what you mentioned later, where torch.stack was returning a tensorExtra object and, for each “element” in this, type(element) was returning tensorExtra.

@KFrank

I have no idea how or why torch.stack() returns a tensorExtra when a
tensorExtra is passed in.

I’m not sure either. I’ve tried testing this out and editing the class definition, but it somehow returns a tensorExtra with the “broken” extra attribute, which is strange, since the __new__ function explicitly requires an extra attribute to be defined for a tensorExtra to be instantiated.

Thank you again for your help! I appreciate it.