Repository: pytorch
Issue: Subclasses with unwrapping `__torch_dispatch__` impls as parameters
Labels: module: nn,triaged,module: __torch_dispatch__,tensor subclass

## Issue description

As pointed out by @another-pjohnson [here](https://github.com/pytorch/pytorch/pull/73459#issuecomment-1121690504), if a tensor subclass implements `__torch_dispatch__` that doesn't re-wrap to return a subclass instance, the Parameter version of the instance will be of the unwrapped type rather than the subclass type.

```python
import torch
from torch.utils._pytree import tree_map

class NonRewrappingTensor(torch.Tensor):
    @staticmethod
    def __new__(
        cls, t: torch.Tensor
    ):
        r = super(NonRewrappingTensor, cls)._make_wrapper_subclass(cls, t.shape, dtype=t.dtype, requires_grad=t.requires_grad, device=t.device)
        return r

    def __init__(self, t) -> None:
        self.tensor: torch.Tensor = t

    __torch_function__ = torch._C._disabled_torch_function_impl

    @classmethod
    def __torch_dispatch__(cls, func, types, args=(), kwargs=None):

        def unwrap(e) -> torch.Tensor:
            if isinstance(e, NonRewrappingTensor):
                t = e.tensor
                return t
            else:
                return e

        r = func(*tree_map(unwrap, args), **tree_map(unwrap, kwargs))
        # Return an unwrapped tensor no longer of original subclass type.
        return r

t = NonRewrappingTensor(torch.randn(3))
param = torch.nn.Parameter(t)
print(type(param))  # torch.Tensor, should be NonRewrappingTensor
```

This happens because `Parameter.__new__` relies on a call to `detach()`, calling the subclass's `__torch_dispatch__` impl.
https://github.com/pytorch/pytorch/blob/3d0e6f169c6745a548ca3e71af9a719ba67dff20/torch/nn/parameter.py#L40

What should be done here to avoid confusion? A few options:
1. Explicitly enforce that any tensor subclass usable as a Parameter be closed under *all* ops (i.e. its `__torch_dispatch__` impl must always re-wrap). I think this is a non-starter, as e.g. the `SparseTensor` example [here](https://github.com/pytorch/pytorch/blob/3c2e0dc6574e1ddc1dda16067200d5b491a18fa7/torch/testing/_internal/common_subclass.py#L99-L154) does not satisfy this and we still want it usable as a Parameter.
2. Explicitly enforce that any tensor subclass usable as a Parameter be closed under some subset of ops. At the very least, this should include `detach()` (I think this should suffice as it's the only op the new parameter mechanism relies on, but open to discussion). This puts the responsibility to specifically handle `detach()` on the subclass writer, but it's possible for us to throw a nice error in `Parameter.__new__` to inform subclass writers if the `detach()` constraint is violated.

```python
    @classmethod
    def __torch_dispatch__(cls, func, types, args=(), kwargs=None):

        def unwrap(e) -> torch.Tensor:
            if isinstance(e, NonRewrappingTensor):
                t = e.tensor
                return t
            else:
                return e

        r = func(*tree_map(unwrap, args), **tree_map(unwrap, kwargs))
        # Specifically handle detach() so this is usable as a Parameter.
        if func is torch.ops.aten.detach.default:
            return NonRewrappingTensor(r)
        return r
```

3. Leave as-is; working as intended.  Are there any valid use cases for using a tensor subclass as parameter and having `detach()` behavior returning the unwrapped type or is this just a problem waiting to happen for subclass writers that need unwrapping behavior in their `__torch_dispatch__` impl?
4. Some other mechanism to avoid invoking `__torch_dispatch__` impl when detaching in `Parameter.__new__`?

cc @albanD @mruberry @jbschlosser @walterddr @kshitij12345 @Chillee @ezyang @zou3519 @samdow