Repository: pytorch
Issue: [JIT] Cannot `jit.export` a `@staticmethod`
Labels: oncall: jit,Stale

### 🐛 Describe the bug

Neither `jit.export` nor `jit.script` work with `@staticmethod` decorated methods:

```python
import torch
from torch import nn, jit, Tensor
from typing import Final

class PositionalEncoder(nn.Module):
    num_dim: Final[int]
    scale: Final[float]
    scales: Tensor

    def __init__(self, num_dim: int, scale: float) -> None:
        super().__init__()
        self.num_dim = num_dim
        self.scale = float(scale)
        scales = self.scale ** (-torch.arange(0, num_dim + 2, 2) / num_dim)
        assert scales[0] == 1.0, "Something went wrong."
        self.register_buffer("scales", scales, persistent=True)

    @jit.export
    def forward(self, t: Tensor) -> Tensor:
        z = torch.einsum("..., d -> ...d", t, self.scales)
        return torch.cat([torch.sin(z), torch.cos(z)], dim=-1)

    @jit.export
    @staticmethod
    def inverse(t: Tensor) -> Tensor:
        return torch.asin(t[..., 0])

encoder = jit.script(PositionalEncoder(5, 1.23))

x = torch.randn(5)
y = encoder(x)
encoder.inverse(y)
```

This results in `AttributeError: 'RecursiveScriptModule' object has no attribute 'inverse'` : 

<details>

```
Traceback (most recent call last):
  File "/home/rscholz/Projects/KIWI/tsdm/bugs/jit_export_staticmethod.py", line 32, in <module>
    encoder.inverse(y)
  File "/home/rscholz/miniconda3/envs/kiwi/lib/python3.9/site-packages/torch/jit/_script.py", line 757, in __getattr__
    return super(RecursiveScriptModule, self).__getattr__(attr)
  File "/home/rscholz/miniconda3/envs/kiwi/lib/python3.9/site-packages/torch/jit/_script.py", line 474, in __getattr__
    return super(ScriptModule, self).__getattr__(attr)
  File "/home/rscholz/miniconda3/envs/kiwi/lib/python3.9/site-packages/torch/nn/modules/module.py", line 1177, in __getattr__
    raise AttributeError("'{}' object has no attribute '{}'".format(
AttributeError: 'RecursiveScriptModule' object has no attribute 'inverse'
```

</details>

If one replaces `jit.export` with `jit.script` we isntead get `AttributeError: 'wrapper_descriptor' object has no attribute '__globals__'`:

<details>

```
Traceback (most recent call last):
  File "/home/rscholz/Projects/KIWI/tsdm/bugs/jit_export_staticmethod.py", line 5, in <module>
    class PositionalEncoder(nn.Module):
  File "/home/rscholz/Projects/KIWI/tsdm/bugs/jit_export_staticmethod.py", line 25, in PositionalEncoder
    def inverse(t: Tensor) -> Tensor:
  File "/home/rscholz/miniconda3/envs/kiwi/lib/python3.9/site-packages/torch/jit/_script.py", line 1318, in script
    return torch.jit._recursive.create_script_class(obj)
  File "/home/rscholz/miniconda3/envs/kiwi/lib/python3.9/site-packages/torch/jit/_recursive.py", line 417, in create_script_class
    rcb = _jit_internal.createResolutionCallbackForClassMethods(type(obj))
  File "/home/rscholz/miniconda3/envs/kiwi/lib/python3.9/site-packages/torch/_jit_internal.py", line 394, in createResolutionCallbackForClassMethods
    captures.update(get_closure(fn))
  File "/home/rscholz/miniconda3/envs/kiwi/lib/python3.9/site-packages/torch/_jit_internal.py", line 168, in get_closure
    captures.update(fn.__globals__)
AttributeError: 'wrapper_descriptor' object has no attribute '__globals__'
```

</details>

### Versions

<details>

```
Collecting environment information...
PyTorch version: 1.10.1
Is debug build: False
CUDA used to build PyTorch: 11.3
ROCM used to build PyTorch: N/A

OS: Ubuntu 20.04.3 LTS (x86_64)
GCC version: (Ubuntu 9.3.0-17ubuntu1~20.04) 9.3.0
Clang version: 10.0.0-4ubuntu1 
CMake version: version 3.16.3
Libc version: glibc-2.31

Python version: 3.9.7 | packaged by conda-forge | (default, Sep 29 2021, 19:20:46)  [GCC 9.4.0] (64-bit runtime)
Python platform: Linux-5.11.0-41-generic-x86_64-with-glibc2.31
Is CUDA available: True
CUDA runtime version: 11.4.152
GPU models and configuration: GPU 0: NVIDIA GeForce RTX 3090
Nvidia driver version: 495.29.05
cuDNN version: Probably one of the following:
/usr/lib/x86_64-linux-gnu/libcudnn.so.8.3.1
/usr/lib/x86_64-linux-gnu/libcudnn_adv_infer.so.8.3.1
/usr/lib/x86_64-linux-gnu/libcudnn_adv_train.so.8.3.1
/usr/lib/x86_64-linux-gnu/libcudnn_cnn_infer.so.8.3.1
/usr/lib/x86_64-linux-gnu/libcudnn_cnn_train.so.8.3.1
/usr/lib/x86_64-linux-gnu/libcudnn_ops_infer.so.8.3.1
/usr/lib/x86_64-linux-gnu/libcudnn_ops_train.so.8.3.1
/usr/local/cuda-11.0/targets/x86_64-linux/lib/libcudnn.so.8.2.1
/usr/local/cuda-11.0/targets/x86_64-linux/lib/libcudnn_adv_infer.so.8.2.1
/usr/local/cuda-11.0/targets/x86_64-linux/lib/libcudnn_adv_train.so.8.2.1
/usr/local/cuda-11.0/targets/x86_64-linux/lib/libcudnn_cnn_infer.so.8.2.1
/usr/local/cuda-11.0/targets/x86_64-linux/lib/libcudnn_cnn_train.so.8.2.1
/usr/local/cuda-11.0/targets/x86_64-linux/lib/libcudnn_ops_infer.so.8.2.1
/usr/local/cuda-11.0/targets/x86_64-linux/lib/libcudnn_ops_train.so.8.2.1
/usr/local/cuda-11.1/targets/x86_64-linux/lib/libcudnn.so.8.2.1
/usr/local/cuda-11.1/targets/x86_64-linux/lib/libcudnn_adv_infer.so.8.2.1
/usr/local/cuda-11.1/targets/x86_64-linux/lib/libcudnn_adv_train.so.8.2.1
/usr/local/cuda-11.1/targets/x86_64-linux/lib/libcudnn_cnn_infer.so.8.2.1
/usr/local/cuda-11.1/targets/x86_64-linux/lib/libcudnn_cnn_train.so.8.2.1
/usr/local/cuda-11.1/targets/x86_64-linux/lib/libcudnn_ops_infer.so.8.2.1
/usr/local/cuda-11.1/targets/x86_64-linux/lib/libcudnn_ops_train.so.8.2.1
/usr/local/cuda-11.2/targets/x86_64-linux/lib/libcudnn.so.8.2.1
/usr/local/cuda-11.2/targets/x86_64-linux/lib/libcudnn_adv_infer.so.8.2.1
/usr/local/cuda-11.2/targets/x86_64-linux/lib/libcudnn_adv_train.so.8.2.1
/usr/local/cuda-11.2/targets/x86_64-linux/lib/libcudnn_cnn_infer.so.8.2.1
/usr/local/cuda-11.2/targets/x86_64-linux/lib/libcudnn_cnn_train.so.8.2.1
/usr/local/cuda-11.2/targets/x86_64-linux/lib/libcudnn_ops_infer.so.8.2.1
/usr/local/cuda-11.2/targets/x86_64-linux/lib/libcudnn_ops_train.so.8.2.1
/usr/local/cuda-11.3/targets/x86_64-linux/lib/libcudnn.so.8.2.1
/usr/local/cuda-11.3/targets/x86_64-linux/lib/libcudnn_adv_infer.so.8.2.1
/usr/local/cuda-11.3/targets/x86_64-linux/lib/libcudnn_adv_train.so.8.2.1
/usr/local/cuda-11.3/targets/x86_64-linux/lib/libcudnn_cnn_infer.so.8.2.1
/usr/local/cuda-11.3/targets/x86_64-linux/lib/libcudnn_cnn_train.so.8.2.1
/usr/local/cuda-11.3/targets/x86_64-linux/lib/libcudnn_ops_infer.so.8.2.1
/usr/local/cuda-11.3/targets/x86_64-linux/lib/libcudnn_ops_train.so.8.2.1
/usr/local/cuda-11.4/targets/x86_64-linux/lib/libcudnn.so.8.2.2
/usr/local/cuda-11.4/targets/x86_64-linux/lib/libcudnn_adv_infer.so.8.2.2
/usr/local/cuda-11.4/targets/x86_64-linux/lib/libcudnn_adv_train.so.8.2.2
/usr/local/cuda-11.4/targets/x86_64-linux/lib/libcudnn_cnn_infer.so.8.2.2
/usr/local/cuda-11.4/targets/x86_64-linux/lib/libcudnn_cnn_train.so.8.2.2
/usr/local/cuda-11.4/targets/x86_64-linux/lib/libcudnn_ops_infer.so.8.2.2
/usr/local/cuda-11.4/targets/x86_64-linux/lib/libcudnn_ops_train.so.8.2.2
HIP runtime version: N/A
MIOpen runtime version: N/A

Versions of relevant libraries:
[pip3] mypy==0.910
[pip3] mypy-extensions==0.4.3
[pip3] numpy==1.21.4
[pip3] numpydoc==1.1.0
[pip3] pytorch-forecasting==0.9.2
[pip3] pytorch-lightning==1.5.6
[pip3] torch==1.10.1
[pip3] torchaudio==0.10.1
[pip3] torchdiffeq==0.2.2
[pip3] torchinfo==1.5.4
[pip3] torchmetrics==0.6.2
[pip3] torchtext==0.11.1
[pip3] torchvision==0.11.2
[conda] blas                      2.112                       mkl    conda-forge
[conda] blas-devel                3.9.0            12_linux64_mkl    conda-forge
[conda] cudatoolkit               11.3.1               ha36c431_9    nvidia
[conda] libblas                   3.9.0            12_linux64_mkl    conda-forge
[conda] libcblas                  3.9.0            12_linux64_mkl    conda-forge
[conda] liblapack                 3.9.0            12_linux64_mkl    conda-forge
[conda] liblapacke                3.9.0            12_linux64_mkl    conda-forge
[conda] mkl                       2021.4.0           h06a4308_640  
[conda] mkl-devel                 2021.4.0           h66538d2_640  
[conda] mkl-include               2021.4.0           h06a4308_640  
[conda] mypy                      0.910              pyhd3eb1b0_0  
[conda] mypy_extensions           0.4.3            py39h06a4308_0  
[conda] numexpr                   2.8.0           mkl_py39h071b83d_0    conda-forge
[conda] numpy                     1.21.4           py39hdbf815f_0    conda-forge
[conda] numpydoc                  1.1.0              pyhd3eb1b0_1  
[conda] pytorch                   1.10.1          py3.9_cuda11.3_cudnn8.2.0_0    pytorch
[conda] pytorch-forecasting       0.9.2              pyhd8ed1ab_0    conda-forge
[conda] pytorch-lightning         1.5.6              pyhd8ed1ab_0    conda-forge
[conda] pytorch-mutex             1.0                        cuda    pytorch
[conda] torchaudio                0.10.1               py39_cu113    pytorch
[conda] torchdiffeq               0.2.2              pyhd8ed1ab_0    conda-forge
[conda] torchinfo                 1.5.4              pyhd8ed1ab_0    conda-forge
[conda] torchmetrics              0.6.2              pyhd8ed1ab_0    conda-forge
[conda] torchtext                 0.11.1                     py39    pytorch
[conda] torchvision               0.11.2               py39_cu113    pytorch
```


</details>