Tensors that share same underlying storage to also share gradient storage
Labels: module: autograd, triaged, needs research, module: python frontend

### 🚀 The feature, motivation and pitch

Hello Pytorch Team! 

I have a question concerning the way gradients are handled, in particular grad attribute. It's a bit hard to explain but in a very short manner I want "views" of a given parameter to have "synced" gradients with the original parameter (shared memory storage). In short the following script tries to mimick what I want to do.

```python
import torch
from torch import nn

def hack_to_sync_grad(param, *view_args, **view_kwargs):
    with torch.no_grad():
        param_view = param.view(*view_args, **view_kwargs)

    # We need sliced tensor to also get the gradient in order to run optimizer on them
    # TODO @thomasw21: It's really had to make sure that our sliced view keeps the same memory space as the original gradient
    def sync_grad(orig_grad):
        assert orig_grad.is_contiguous()
        if param_view.grad is None:
            # The gradient was reset to None, we need to reset the coalesced_tensor.grad as well
            param.grad = None

        if param.grad is None:
            param_view.grad = orig_grad.view(*view_args, **view_kwargs)
        else:
            assert param_view.grad is not None
            # TODO @thomasw21: How do I check that it's still the same memory allocation?
            # One way is to re-assign the existing pointer
            param_view.grad = param.grad.view(*view_args, **view_kwargs)
        return orig_grad

    param.register_hook(sync_grad)
    return param_view

def main():
    model = nn.Linear(2,3)

    weight_view = model.weight.view(-1) # This fails as it doesn't propagate gradients to this view
    # weight_view = hack_to_sync_grad(model.weight, -1) # This is fairly hacky since it relies on the backward hook to sync

    random_input = torch.randn(5,2)
    model(random_input).sum().backward()

    for name, param in model.named_parameters():
        assert param.grad is not None

    # My goal is such that a view of a parameters should also be able to get the "sharded" gradient
    assert weight_view.grad is not None
    assert weight_view.grad.untyped_storage().data_ptr() == model.weight.grad.untyped_storage().data_ptr()

if __name__ == "__main__":
    main()
```

In this example I'm using `view` but I'd like to support things like slices as well. So very concretely what I'm interested in is customising a grad getter. This would allow one to have a syncing mechanism that doesn't rely on the backward hook for syncing.

Solution found as of now:
 - Creating a new property of a new class:
 ```python
 class ViewTensor(torch.Tensor):
    """We want view tensor"""
    def __new__(cls, data, *view_args, **view_kwargs):
        new_tensor = super().__new__(cls, data.view(*view_args, **view_kwargs))
        new_tensor.orig_data = data
        new_tensor.view_args = view_args
        new_tensor.view_kwargs = view_kwargs
        return new_tensor

    @property
    def grad(self):
        return self.orig_data.grad.view(*self.view_args, **self.view_kwargs)

    @grad.setter
    def grad(self, value):
        self.orig_data.grad = value.view_as(self.orig_data)

    @grad.deleter
    def grad(self):
        del self.orig_data.grad

def use_new_tensor_class(param, *view_args, **view_kwargs):
    with torch.no_grad():
        param_view = ViewTensor(param, *view_args, **view_kwargs)
    return param_view
 ```
 The only issue is that property is not something we can set on an instance, hence the need to define a new class. Maybe this solution can generalize by allowing for torch.Tensor to pass grad getter/setter/deleter
 
 cc @ezyang @albanD @zou3519 @gqchen @pearu @nikitaved @soulitzer @Lezcano @Varal7 with whom I initially discussed

### Alternatives

_No response_

### Additional context

_No response_