Repository: pytorch
Issue: Certain torch functions are not handled by torch func wrapper
Labels: triaged,module: nestedtensor

### 🐛 Describe the bug

When using custom Tensor class, we register functions in a function registry and call them in `__torch_function__` of the custom class.

This behavior is described in [docs](https://pytorch.org/docs/stable/notes/extending.html#extending-torch).

However, it only works for certain functions, e.g., (`torch.mean`), but when I trying to call with `torch.cat`, it complains:

```
TypeError: cat(): argument 'tensors' (position 1) must be tuple of Tensors, not NestedTensor
```

I have the following [code](https://github.com/ZhiyuanChen/DanLing/blob/8230fe6bdcace72c00d34a77b7675fa084b42193/danling/tensors/nested_tensor.py#L552-L559):

```python
    @classmethod
    def __torch_function__(cls, func, types, args=(), kwargs=None):
        if kwargs is None:
            kwargs = {}
        if func not in NestedTensorFunc or not all(issubclass(t, (torch.Tensor, NestedTensor)) for t in types):
            args = [a.tensor if hasattr(a, "tensor") else a for a in args]
            return func(*args, **kwargs)
        return NestedTensorFunc[func](*args, **kwargs)
```

And the [function registered](https://github.com/ZhiyuanChen/DanLing/blob/8230fe6bdcace72c00d34a77b7675fa084b42193/danling/tensors/nested_tensor.py#L632-L661) is :

```python
@NestedTensorFunc.implement(torch.mean)  # pylint: disable=E1101
def mean(
    input,  # pylint: disable=W0622
    dim: int | None = None,
    keepdim: bool = False,
    *,
    dtype: torch.dtype | None = None,
):
    return input.mean(dim=dim, keepdim=keepdim, dtype=dtype)


@NestedTensorFunc.implement(torch.cat)  # pylint: disable=E1101
def cat(tensors, dim: int = 0):
    if dim != 0:
        raise NotImplementedError(f"NestedTensor only supports cat when dim=0, but got {dim}")
    return NestedTensor([t for tensor in tensors for t in tensor.storage], tensors[0]._state())  # pylint: disable=W0212


@NestedTensorFunc.implement(torch.stack)  # pylint: disable=E1101
def stack(tensors, dim: int = 0):
    raise NotImplementedError("NestedTensor does not support stack as of now")


@NestedTensorFunc.implement(torch.isin)  # pylint: disable=E1101
def isin(elements, test_elements, *, assume_unique: bool = False, invert: bool = False):
    if isinstance(elements, NestedTensor):
        elements = elements.tensor
    if isinstance(test_elements, NestedTensor):
        test_elements = test_elements.tensor
    return torch.isin(elements, test_elements, assume_unique=assume_unique, invert=invert)  # pylint: disable=E1101
```

### Versions

PyTorch version: 2.0.1
Is debug build: False
CUDA used to build PyTorch: 11.8
ROCM used to build PyTorch: N/A

OS: Ubuntu 23.04 (x86_64)
GCC version: (Ubuntu 12.3.0-1ubuntu1~23.04) 12.3.0
Clang version: 16.0.0 (1~exp5ubuntu3)
CMake version: version 3.25.1
Libc version: glibc-2.37

Python version: 3.10.12 | packaged by conda-forge | (main, Jun 23 2023, 22:40:32) [GCC 12.3.0] (64-bit runtime)
Python platform: Linux-6.2.0-27-generic-x86_64-with-glibc2.37
Is CUDA available: True
CUDA runtime version: 11.8.89
CUDA_MODULE_LOADING set to: LAZY
GPU models and configuration:
GPU 0: NVIDIA A100-SXM4-80GB
GPU 1: NVIDIA A100-SXM4-80GB
GPU 2: NVIDIA A100-SXM4-80GB
GPU 3: NVIDIA A100-SXM4-80GB
GPU 4: NVIDIA A100-SXM4-80GB
GPU 5: NVIDIA A100-SXM4-80GB
GPU 6: NVIDIA A100-SXM4-80GB
GPU 7: NVIDIA A100-SXM4-80GB

Nvidia driver version: 535.54.03
cuDNN version: Could not collect
HIP runtime version: N/A
MIOpen runtime version: N/A
Is XNNPACK available: True

CPU:
Architecture:                    x86_64
CPU op-mode(s):                  32-bit, 64-bit
Address sizes:                   46 bits physical, 57 bits virtual
Byte Order:                      Little Endian
CPU(s):                          112
On-line CPU(s) list:             0-111
Vendor ID:                       GenuineIntel
Model name:                      Intel(R) Xeon(R) Platinum 8336C CPU @ 2.30GHz
CPU family:                      6
Model:                           106
Thread(s) per core:              2
Core(s) per socket:              28
Socket(s):                       2
Stepping:                        6
BogoMIPS:                        4600.07
Flags:                           fpu vme de pse tsc msr pae mce cx8 apic sep mtrr pge mca cmov pat pse36 clflush mmx fxsr sse sse2 ss ht syscall nx pdpe1gb rdtscp lm constant_tsc rep_good nopl xtopology cpuid tsc_known_freq pni pclmulqdq ssse3 fma cx16 pcid sse4_1 sse4_2 x2apic movbe popcnt tsc_deadline_timer aes xsave avx f16c rdrand hypervisor lahf_lm abm 3dnowprefetch cpuid_fault invpcid_single ssbd ibrs ibpb stibp ibrs_enhanced fsgsbase tsc_adjust bmi1 avx2 smep bmi2 erms invpcid avx512f avx512dq rdseed adx smap avx512ifma clflushopt clwb avx512cd sha_ni avx512bw avx512vl xsaveopt xsavec xgetbv1 xsaves wbnoinvd arat avx512vbmi umip pku ospke avx512_vbmi2 gfni vaes vpclmulqdq avx512_vnni avx512_bitalg avx512_vpopcntdq la57 rdpid md_clear arch_capabilities
Hypervisor vendor:               KVM
Virtualization type:             full
L1d cache:                       2.6 MiB (56 instances)
L1i cache:                       1.8 MiB (56 instances)
L2 cache:                        70 MiB (56 instances)
L3 cache:                        54 MiB (1 instance)
NUMA node(s):                    2
NUMA node0 CPU(s):               0-55
NUMA node1 CPU(s):               56-111
Vulnerability Itlb multihit:     Not affected
Vulnerability L1tf:              Not affected
Vulnerability Mds:               Not affected
Vulnerability Meltdown:          Not affected
Vulnerability Mmio stale data:   Vulnerable: Clear CPU buffers attempted, no microcode; SMT Host state unknown
Vulnerability Retbleed:          Not affected
Vulnerability Spec store bypass: Mitigation; Speculative Store Bypass disabled via prctl
Vulnerability Spectre v1:        Mitigation; usercopy/swapgs barriers and __user pointer sanitization
Vulnerability Spectre v2:        Mitigation; Enhanced IBRS, IBPB conditional, RSB filling, PBRSB-eIBRS SW sequence
Vulnerability Srbds:             Not affected
Vulnerability Tsx async abort:   Not affected

Versions of relevant libraries:
[pip3] flake8==6.0.0
[pip3] mypy==0.981
[pip3] mypy-extensions==0.4.3
[pip3] numpy==1.25.2
[pip3] torch==2.0.1
[pip3] torchaudio==2.0.2
[pip3] torcheval==0.0.6
[pip3] torchmetrics==1.0.3
[pip3] torchtnt==0.2.0
[pip3] torchvision==0.15.2
[pip3] triton==2.0.0
[conda] blas                      1.0                         mkl    defaults
[conda] ffmpeg                    4.3                  hf484d3e_0    pytorch
[conda] mkl                       2023.1.0         h213fc3f_46343    defaults
[conda] mkl-service               2.4.0           py310h5eee18b_1    defaults
[conda] mkl_fft                   1.3.6           py310h1128e8f_1    defaults
[conda] mkl_random                1.2.2           py310h1128e8f_1    defaults
[conda] numpy                     1.25.2          py310h5f9d8c6_0    defaults
[conda] numpy-base                1.25.2          py310hb5e798b_0    defaults
[conda] pytorch                   2.0.1           py3.10_cuda11.8_cudnn8.7.0_0    pytorch
[conda] pytorch-cuda              11.8                 h7e8668a_5    pytorch
[conda] pytorch-mutex             1.0                        cuda    pytorch
[conda] torchaudio                2.0.2               py310_cu118    pytorch
[conda] torcheval                 0.0.6                    pypi_0    pypi
[conda] torchmetrics              1.0.3                    pypi_0    pypi
[conda] torchtnt                  0.2.0                    pypi_0    pypi
[conda] torchtriton               2.0.0                     py310    pytorch
[conda] torchvision               0.15.2              py310_cu118    pytorch

cc @cpuhrsch @jbschlosser @bhosmer @drisspg