Make subclass Tensor has grad attributes


I want to subclass the Tensor and make it has .grad attribute during backward.

As .grad is only available for leaf tensor, I made the following simple demo but found the behavior is not consistent with original torch.Tensor. There will be AliasBackward introduced in the subclass tensor, but I want to remove it to make it looks like original torch.Tensor.

Current it has the behavior like this:

import torch

class TensorFoo(torch.Tensor): pass
tensor = TensorFoo([1,2,3,4]).cuda().requires_grad_()
print(tensor.is_leaf, tensor.requires_grad)
tensor = torch.Tensor([1,2,3,4]).cuda().requires_grad_()
print(tensor.is_leaf, tensor.requires_grad)

The output will be

False True
True True

Have any ideas about this? Thanks in advance.


I can reproduce this on 1.9.0 but it seems to be fixed on master.

# !pip install --pre torch -f
print(torch.version.__version__) # '1.10.0.dev20210723+cu111'
class TensorFoo(torch.Tensor): pass
tensor = TensorFoo([1,2,3,4]).requires_grad_()
print(tensor.is_leaf, tensor.requires_grad) # True True

As a workaround, you can use tensor._base