Repository: pytorch
Issue: Proposal: Add __tensor_wrap__ method similar to numpy __array_wrap__
Labels: module: internals,triaged,module: numpy

## 🚀 Feature
Add the ability to subclass torch.Tensor so that you can use `isinstance(tensor, CustomTensor)` to determine the "context" of the data stored in the subclass object. I think #9515 describes the need for **dtype** promotion logic, but this proposal is for **python type** promotion logic. 

The __array__wrap__ method already exists in numpy and ensures that the following ufunc preserves the custom tensor type:
```python
custom_tensor = CustomTensor(torch.tensor([1, 2, 3, 4])
new_tensor = np.add(custom_tensor, 1.0) -> CustomTensor
new_tensor = np.add(1.0, custom_tensor) -> CustomTensor (Hmmm...)
isinstance(new_tensor, CustomTensor) #  returns True
```

**The __tensor_wrap__ method does not exist**. As a result, the custom type information is easily lost:
custom_tensor = CustomTensor(torch.tensor([1, 2, 3, 4])
new_tensor = custom_tensor.add(1.0) -> Tensor
isinstance(new_tensor, CustomTensor) #  returns False

## Motivation

The torch.Tensor class provides a dynamically typed wrapper of statically typed Torch Classes (FloatTensor, DoubleTensor, HalfTensor, etc) using the numpy dtype. While this should save us from ever wanting to subclass the Tensor class, it would be nice to use the subclass of a Tensor to provide custom methods specific to the context of the data. 

For example: 
```python
class TimeDomainTensor(torch.Tensor)

    def __repr__()
        return 'Time Information:\n' + super(Parameter, self).__repr__()
```
```python
class FrequencyDomainTensor(torch.Tensor)
   
    def __repr__()
        return 'Frequency Information:\n' + super(Parameter, self).__repr__()
```
```python
class SurfaceTensor(torch.Tensor)
   
    def __repr__()
        return 'Surface Information:\n' + super(Parameter, self).__repr__()
```

## Pitch

Custom subclassing is already utilized by the torch.nn.Parameter Class, so most of the heavy lifting is already done. Note that performing a math operation on a Parameter also downcasts the Parameter to a Tensor.

## Alternatives

Numpy provides the following methods for customizing the output behaviour of a ufunc.

1. def __array_wrap__(self, array):
    A subclass can override what happens when executing numpy ufuncs **after** the ufunc is called.

2. def __array_ufunc__(ufunc, method, *inputs, **kwargs):
    A subclass can override what happens when executing numpy ufuncs **before and after** the ufunc is called. New in version 1.13.
 
3. def __array_finalize__(self, obj): 
    The ability to insert any kind of custom attributes, events, io, etc.  This may not work on GPU because this would open up the ability to define highly customized data containers.
        a. explicit constructor call (obj = MySubClass(params)).
        b. View casting
        c. Creating new from template

4. Do Nothing:
   Revert to functional implementation. Simply replace all custom_tensor.method(*args, **kwargs) with method(custom_tensor, *args, **kwargs).

The __array_wrap__ method is sufficient for this problem because we want to cast the tensor **after** the function has been performed. The __array_ufunc__ would also work but could result in compatibility  problems with older versions of numpy. The __array_finalize method may add too much customization that could create a wide range of gotchas and incompatibility between CPU and GPU.

## Example

Sample code that partially works for add(custom_tensor, other):

```python
from collections import OrderedDict
import torch as th


def _rebuild_custom_tensor(data, requires_grad, backward_hooks):
    param = CustomTensor(data, requires_grad)
    # NB: This line exists only for backwards compatibility; the
    # general expectation is that backward_hooks is an empty
    # OrderedDict.  See Note [Don't serialize hooks]
    param._backward_hooks = backward_hooks

    return param

class CustomTensor(th.Tensor):
    r"""A kind of Tensor that has time, freq as its inner-most dimensions.

    Arguments:
        data (Tensor): parameter tensor.
        requires_grad (bool, optional): if the parameter requires gradient. See
            :ref:`excluding-subgraphs` for more details. Default: `True`
    """

    def __new__(cls, data=None, requires_grad=False):
        if data is None:
            data = th.Tensor()
        return th.Tensor._make_subclass(cls, data, requires_grad)

    def __deepcopy__(self, memo):
        if id(self) in memo:
            return memo[id(self)]
        else:
            result = type(self)(self.data.clone(), self.requires_grad)
            memo[id(self)] = result
            return result

    def __repr__(self):
        return 'Custom Message:\n' + super(CustomTensor, self).__repr__()

    def __reduce_ex__(self, proto):
        # See Note [Don't serialize hooks]
        return (
            _rebuild_custom_tensor,
            (self.data, self.requires_grad, OrderedDict())
        )

    def add(self, value):
        tensor = super().add(value)
        return self.__tensor_wrap__(tensor)

    def __add__(self, other):
        tensor = super().__add__(other)
        return self.__tensor_wrap__(tensor)

    # Wrap Numpy array again in a suitable tensor when done, to support e.g.
    # `numpy.sin(tensor) -> tensor` or `numpy.greater(tensor, 0) -> ByteTensor`
    def __array_wrap__(self, array):
        if array.dtype == bool:
            # Workaround, torch has no built-in bool tensor
            array = array.astype('uint8')
            return th.from_numpy(array)
        else:
            return th.Tensor._make_subclass(CustomTensor, th.from_numpy(array), self.requires_grad)

    # Same as __array_wrap__ but for Tensors
    def __tensor_wrap__(self, tensor):
        return th.Tensor._make_subclass(CustomTensor, tensor, self.requires_grad)
```

Sample Results:
```python
import torch
from custom import CustomTensor
x = torch.tensor([1.0, 2.0, 3.0], dtype=torch.double)
custom_x = CustomTensor(x)
custom_x + 1   # -> CustomTensor
1 + custom_x   # -> **Tensor** :(
torch.add(1, custom_tensor) -> **Tensor** :(
torch.add(custom_tensor, 1) -> **Tensor** :(
```

## Conclusions
The above example demonstrates how numpy goes beyond binary operator overloading to ensure that all operations return the highest derived subclass of the ndarray. This is not how pytorch is organized. I think #9515 describes the need for dtype promotion logic, but this proposal is for python type promotion logic.