Repository: pytorch
Issue: using Tensor subclass between vmap layers
Labels: triaged,module: __torch_dispatch__,tensor subclass,module: functorch

### 🐛 Describe the bug

Trying to use Tensor subclass between vmap layers:
```python
import torch
import functorch
from torch.utils._pytree import tree_map

class CustomTensor(torch.Tensor):
    def __new__(cls, elem):
        self = torch.Tensor._make_wrapper_subclass(
            cls, elem.size(),
            dtype=elem.dtype, layout=elem.layout,
            device=elem.device, requires_grad=elem.requires_grad,
            strides=elem.stride(), storage_offset=elem.storage_offset())
        self.elem = elem
        return self
    def __repr__(self):
        return f'{type(self).__name__}({repr(self.elem)})'
    @classmethod
    def __torch_dispatch__(cls, func, types, args=(), kargs=None):
        print('dispatch', func, types)
        if kargs is None:
            kargs = {}
        memo = {}
        unwrap = lambda x : x.elem if type(x) is cls else x
        wrap = lambda x : cls(x) if torch.overrides.is_tensor_like(x) else x

        args = tree_map(unwrap, args)
        kargs = tree_map(unwrap, kargs)

        with torch.overrides.enable_reentrant_dispatch():
            ans = func(*args, **kargs)
        ans = tree_map(wrap, ans)
        return ans

def f(x):
    def g(y):
        return y + y
    a = functorch.vmap(g)(CustomTensor(x))
    return a.elem
functorch.vmap(f)(torch.tensor([[1., 2.], [3., 4.]]))
```
This outputs the following:
```
dispatch aten.add.Tensor (<class '__main__.CustomTensor'>,)
Traceback (most recent call last):
  File "/home/alan/test.py", line 38, in <module>
    functorch.vmap(f)(torch.tensor([[1., 2.], [3., 4.]]))
  File "/home/alan/env/lib/python3.10/site-packages/torch/_functorch/vmap.py", line 420, in wrapped
    return _flat_vmap(
  File "/home/alan/env/lib/python3.10/site-packages/torch/_functorch/vmap.py", line 39, in fn
    return f(*args, **kwargs)
  File "/home/alan/env/lib/python3.10/site-packages/torch/_functorch/vmap.py", line 605, in _flat_vmap
    batched_outputs = func(*batched_inputs, **kwargs)
  File "/home/alan/test.py", line 36, in f
    a = functorch.vmap(g)(CustomTensor(x))
  File "/home/alan/env/lib/python3.10/site-packages/torch/_functorch/vmap.py", line 420, in wrapped
    return _flat_vmap(
  File "/home/alan/env/lib/python3.10/site-packages/torch/_functorch/vmap.py", line 39, in fn
    return f(*args, **kwargs)
  File "/home/alan/env/lib/python3.10/site-packages/torch/_functorch/vmap.py", line 605, in _flat_vmap
    batched_outputs = func(*batched_inputs, **kwargs)
  File "/home/alan/test.py", line 35, in g
    return y + y
  File "/home/alan/test.py", line 29, in __torch_dispatch__
    ans = func(*args, **kargs)
  File "/home/alan/env/lib/python3.10/site-packages/torch/_ops.py", line 284, in __call__
    return self._op(*args, **kwargs or {})
RuntimeError: Either your tensor may have escaped from inside a function being vmapped and this is a user error (see https://pytorch.org/functorch/stable/ux_limitations.html), or there is an internal functorch error in `gen_vmap_plumbing` Please file an issue if it looks like the latter
```
without enable_reentrant_dispatch, it also throws:
```
RuntimeError: NYI: querying is_contiguous inside of vmap for memory_format other than torch.contiguous_format
```

### Versions

PyTorch version: 2.0.0.dev20221227+cu116
Is debug build: False
CUDA used to build PyTorch: 11.6
ROCM used to build PyTorch: N/A

OS: Ubuntu 20.04.5 LTS (x86_64)
GCC version: (Ubuntu 9.4.0-1ubuntu1~20.04.1) 9.4.0
Clang version: 10.0.0-4ubuntu1
CMake version: version 3.25.0
Libc version: glibc-2.31

Python version: 3.10.7 (main, Sep  7 2022, 15:22:19) [GCC 9.4.0] (64-bit runtime)
Python platform: Linux-5.10.102.1-microsoft-standard-WSL2-x86_64-with-glibc2.31
Is CUDA available: True
CUDA runtime version: 11.8.89
CUDA_MODULE_LOADING set to: LAZY
GPU models and configuration: GPU 0: NVIDIA GeForce GTX 1650
Nvidia driver version: 512.72
cuDNN version: Probably one of the following:
/usr/lib/x86_64-linux-gnu/libcudnn.so.8.6.0
/usr/lib/x86_64-linux-gnu/libcudnn_adv_infer.so.8.6.0
/usr/lib/x86_64-linux-gnu/libcudnn_adv_train.so.8.6.0
/usr/lib/x86_64-linux-gnu/libcudnn_cnn_infer.so.8.6.0
/usr/lib/x86_64-linux-gnu/libcudnn_cnn_train.so.8.6.0
/usr/lib/x86_64-linux-gnu/libcudnn_ops_infer.so.8.6.0
/usr/lib/x86_64-linux-gnu/libcudnn_ops_train.so.8.6.0
HIP runtime version: N/A
MIOpen runtime version: N/A
Is XNNPACK available: True

Versions of relevant libraries:
[pip3] numpy==1.24.1
[pip3] torch==2.0.0.dev20221227+cu116
[pip3] torchaudio==2.0.0.dev20221227+cu116
[pip3] torchtriton==2.0.0+0d7e753227
[pip3] torchvision==0.15.0.dev20221227+cu116
[conda] Could not collect

cc @Chillee @ezyang @zou3519 @albanD @samdow @msaroufim @soumith