[TorchScript] Better support for non-nn.Module classes defined in different file

[davidberard98]

Stack from ghstack (oldest at bottom):

• -> [TorchScript] Better support for non-nn.Module classes defined in different file #125950

Motivating Issue:

We had this example:

# a.py
import torch
from math import sqrt

# This file is used by test_class_type_closure

class MySqrtClass:
    value: float

    def __init__(self):
        self.value = 4.0

    def get_sqrt_inverse(self) -> float:
        return sqrt(1.0 / self.value)

# b.py
import torch
from a import MySqrtClass

torch.jit.script(MySqrtClass)

This would fail to script:

Traceback (most recent call last):
  File "/data/users/dberard/scripts/ts_use_sqrt.py", line 10, in <module>
    torch.jit.script(ts_sqrt.MyConfig)
  File "/data/users/dberard/pytorch/torch/jit/_script.py", line 1378, in script
    _compile_and_register_class(obj, _rcb, qualified_name)
  File "/data/users/dberard/pytorch/torch/jit/_recursive.py", line 73, in _compile_and_register_class
    script_class = torch._C._jit_script_class_compile(
                   ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
RuntimeError:
undefined value sqrt:
  File "/data/users/dberard/scripts/ts_sqrt.py", line 17
    def get_sqrt_inverse(self) -> float:
        return sqrt(1 / self.value)

Fix details:
In TorchScript, we use "resolution callbacks" - python callables that will map a string to an object. In this case, the resolution callback was unable to identify that sqrt was actually math.sqrt.

We have multiple types of resolution callbacks: two are createResolutionCallbackFromFrame and createResolutionCallbackFromClosure.

In this case, we were using createResolutionCallbackFromFrame. But this doesn't work if jit.script is called from a different file than the class definition.

Fix: createResolutionCallbackFromClosure for each of the methods on a class (which is similar to what we do for an nn.Module). Pass this into the C++ function so that we can use the individual RCBs for each of the methods. If a method doesn't have an RCB for some reason, then fallback to the old RCB method where we look at the current frame.

Differential Revision: D57224315

[pytorch-bot]

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/125950

• 📄 Preview Python docs built from this PR
• 📄 Preview C++ docs built from this PR
• ❓ Need help or want to give feedback on the CI? Visit the bot commands wiki or our office hours

Note: Links to docs will display an error until the docs builds have been completed.

❌ 13 New Failures

As of commit 2bd7446 with merge base 946b96f ():

NEW FAILURES - The following jobs have failed:

• pull / linux-focal-cuda12.1-py3.10-gcc9 / test (default, 2, 5, linux.4xlarge.nvidia.gpu) (gh)
  test_jit.py::TestRecursiveScript::test_attributes
• pull / linux-focal-cuda12.1-py3.10-gcc9-sm86 / test (default, 2, 5, linux.g5.4xlarge.nvidia.gpu) (gh)
  test_jit.py::TestRecursiveScript::test_attributes
• pull / linux-focal-py3.11-clang10 / test (crossref, 2, 2, linux.2xlarge) (gh)
  test_jit.py::TestRecursiveScript::test_attributes
• pull / linux-focal-py3.11-clang10 / test (default, 2, 3, linux.2xlarge) (gh)
  test_jit.py::TestRecursiveScript::test_attributes
• pull / linux-focal-py3.11-clang10 / test (dynamo, 3, 3, linux.2xlarge) (gh)
  test_jit.py::TestRecursiveScript::test_attributes
• pull / linux-focal-py3.12-clang10 / test (default, 2, 3, linux.2xlarge) (gh)
  test_jit.py::TestRecursiveScript::test_attributes
• pull / linux-focal-py3.12-clang10 / test (dynamo, 3, 3, linux.2xlarge) (gh)
  test_jit.py::TestRecursiveScript::test_attributes
• pull / linux-focal-py3.8-clang10 / test (crossref, 2, 2, linux.2xlarge) (gh)
  test_jit.py::TestRecursiveScript::test_attributes
• pull / linux-focal-py3.8-clang10 / test (default, 2, 3, linux.2xlarge) (gh)
  test_jit.py::TestRecursiveScript::test_attributes
• pull / linux-focal-py3.8-clang10 / test (dynamo, 2, 3, linux.2xlarge) (gh)
  test_jit.py::TestRecursiveScript::test_attributes
• pull / linux-jammy-py3.10-clang15-asan / test (default, 2, 6, linux.4xlarge) (gh)
  test_jit.py::TestRecursiveScript::test_attributes
• pull / linux-jammy-py3.8-gcc11 / test (default, 2, 3, linux.2xlarge) (gh)
  test_jit.py::TestRecursiveScript::test_attributes
• pull / linux-jammy-py3.8-gcc11 / test (jit_legacy, 1, 1, linux.2xlarge) (gh)
  test_jit_legacy.py::TestRecursiveScript::test_attributes

This comment was automatically generated by Dr. CI and updates every 15 minutes.

[davidberard98]

@davidberard98 has imported this pull request. If you are a Meta employee, you can view this diff on Phabricator.