Repository: pytorch
Issue: Cannot script when a static method or cross module function calls custom op
Labels: oncall: jit,module: cpp-extensions,internals,module: torchbind

## 🐛 Bug

<!-- A clear and concise description of what the bug is. -->

## To Reproduce

```python
import tempfile

import torch
import torch.utils.cpp_extension

from utils import test

CUSTOM_OP_CPP_SRC = '''#include <torch/script.h>

TORCH_LIBRARY(test, m) {
  m.def("test", [](const std::string& str){
      return str;
  });
}
'''

def setup_module(module):
    with tempfile.NamedTemporaryFile(mode="w", suffix=".cpp", delete=False) as f:
        f.write(CUSTOM_OP_CPP_SRC)
        torch.utils.cpp_extension.load(
            "test",
            [f.name],
            is_python_module=False,
        )

def test_script_calling_static_method_without_custom_op():
    class ClassWithStaticMethod:
        @staticmethod
        def test(some_str: str):
            return some_str

    class TestModule(torch.nn.Module):
        def __init__(self):
            super().__init__()
        
        def forward(self, some_str: str):
            return ClassWithStaticMethod.test(some_str)

    torch.jit.script(TestModule())

def test_script_calling_static_method_with_custom_op():
    class ClassWithStaticMethod:
        @staticmethod
        def test(some_str: str):
            return torch.ops.test.test(some_str)

    class TestModule(torch.nn.Module):
        def __init__(self):
            super().__init__()
        
        def forward(self, some_str: str):
            return ClassWithStaticMethod.test(some_str)

    torch.jit.script(TestModule())

def test_script_calling_cross_module_function_with_custom_op():
    class TestModule(torch.nn.Module):
        def __init__(self):
            super().__init__()
        
        def forward(self, some_str: str):
            return test(some_str)

    torch.jit.script(TestModule())
```

utils.py
```python
import tempfile

import torch
import torch.utils.cpp_extension

CUSTOM_OP_CPP_SRC = '''#include <torch/script.h>

TORCH_LIBRARY(test_cross_module, m) {
  m.def("test_cross_module", [](const std::string& str){
      return str;
  });
}
'''

with tempfile.NamedTemporaryFile(mode="w", suffix=".cpp", delete=False) as f:
    f.write(CUSTOM_OP_CPP_SRC)
    torch.utils.cpp_extension.load(
        "test_cross_module",
        [f.name],
        is_python_module=False,
    )

def test(some_str: str):
    return torch.ops.test_cross_module.test_cross_module(some_str)

```

pytest command: `python -m pytest test_script.py --forked`

`test_script_calling_static_method_without_custom_op` passes

```
=========================================================================================== short test summary info ============================================================================================
FAILED test_script.py::test_script_calling_static_method_with_custom_op
FAILED test_script.py::test_script_calling_function_with_custom_op
```

See stack trace from pytest in comment.

## Expected behavior

The tests should pass.

## Environment
```
(base)  robmeng@Robs-MacBook-Pro  ~/workspaces/pytorch-static-method  python collect_env.py
Collecting environment information...
PyTorch version: 1.6.0
Is debug build: False
CUDA used to build PyTorch: None
ROCM used to build PyTorch: N/A

OS: Mac OSX 10.15.5 (x86_64)
GCC version: Could not collect
Clang version: 11.0.3 (clang-1103.0.32.62)
CMake version: version 3.14.0

Python version: 3.7 (64-bit runtime)
Is CUDA available: False
CUDA runtime version: No CUDA
GPU models and configuration: No CUDA
Nvidia driver version: No CUDA
cuDNN version: No CUDA
HIP runtime version: N/A
MIOpen runtime version: N/A

Versions of relevant libraries:
[pip3] numpy==1.18.4
[pip3] torch2trt==0.0.3
[conda] blas                      1.0                         mkl
[conda] mkl                       2019.4                      233
[conda] mkl-include               2020.2                      260
[conda] mkl-service               2.3.0            py37hfbe908c_0
[conda] mkl_fft                   1.0.15           py37h5e564d8_0
[conda] mkl_random                1.1.0            py37ha771720_0
[conda] numpy                     1.18.1           py37h7241aed_0
[conda] numpy-base                1.18.1           py37h6575580_1
[conda] numpydoc                  0.9.2                      py_0
[conda] torch                     1.6.0                    pypi_0    pypi
```

## Additional context

I think the same problem exists when trying to pass a custom class to a python function. BUT, when passing a custom class to a custom op or custom class method script works as expected.

<!-- Add any other context about the problem here. -->


cc @ezyang @bhosmer @smessmer @ljk53 @bdhirsh @ailzhang @yf225 @glaringlee @gmagogsfm