What's the difference between torch.Tensor._make_subclass and torch.Tensor._make_wrapper_subclass

1

Both of torch.Tensor._make_subclass and torch.Tensor._make_wrapper_subclass are used in Subclassing torch.Tensor such as subclass_zoo/quantized_tensor.py at main · albanD/subclass_zoo · GitHub and pytorch/torch/_subclasses/fake_tensor.py at main · pytorch/pytorch · GitHub
What’s the difference between them?
What’s more, is there any document illustrating how to subclassing torch.Tensor. I find there are many discussions such as Subclassing torch.Tensor - PyTorch Forums and examples GitHub - albanD/subclass_zoo, but I can’t find any document about subclassing torch.Tensor in detail. The document in Extending PyTorch — PyTorch 2.1 documentation is not qualified for users to define subclasses of torch.Tensor. Here are some questions about subclassing torch.Tensor:

  1. In the __new__ method, Sometimes _make_wrapper_subclass is used, sometimes _make_subclass is used, and sometimes neither is used, which confused me a lot.

  2. How should the __new__ and __init__ methods be defined? In other words, what are the responsibilities of __new__ and __init__?

  3. When should we use __torch_function__, and when should we use __torch_dispatch__? Should we prioritize using __torch_dispatch__ according to What (and Why) is __torch_dispatch__?

5 Likes
2

Hey!

The Extending PyTorch doc does talk about subclasses of Tensor a bit below from where you showed: Extending PyTorch — PyTorch main documentation

  1. These two functions do quite different things. The main difference is that when you do _make_subclass(), the current object is a honest to goodness Tensor with data in its storage and everything. When you do _make_wrapper_subclass(), the current object has no data and it is expected that some field on the Tensor will be another Tensor (hence the outer one being called wrapper) that contains real data.
  2. This is the same as any Python class: new creates a new instance while init initializes it. It makes very little difference unless you’re considering serialization of objects. General Python doc does cover that in details though.
  3. It depends on what you want to do. They are different tools for different jobs. The Extending doc tries to give details, this poster from the PyTorch conference tries to answer that question as well: BackToPython PTC 2022 Poster - Google Slides

Hope this helps

2 Likes