# Owner(s): ["oncall: jit"] import os import sys import torch from torch.testing._internal.jit_utils import JitTestCase, make_global from torch.jit._monkeytype_config import _IS_MONKEYTYPE_INSTALLED from typing import List, Dict, Tuple, Any, Optional, NamedTuple # noqa: F401 # Make the helper files in test/ importable pytorch_test_dir = os.path.dirname(os.path.dirname(os.path.realpath(__file__))) sys.path.append(pytorch_test_dir) if not _IS_MONKEYTYPE_INSTALLED: print("monkeytype is not installed. Skipping tests for Profile-Directed Typing", file=sys.stderr) JitTestCase = object # type: ignore[misc, assignment] # noqa: F811 if __name__ == "__main__": raise RuntimeError( "This test file is not meant to be run directly, use:\n\n" "\tpython test/test_jit.py TESTNAME\n\n" "instead." ) class TestPDT(JitTestCase): """ A suite of tests for profile directed typing in TorchScript. """ def test_nn_module(self): class TestPDTModel(torch.nn.Module): def __init__(self): super().__init__() def forward(self, x) -> Any: if isinstance(x, int): return x + 1 elif isinstance(x, float): return x - 1 else: return x make_global(TestPDTModel) pdt_model = TestPDTModel() inp: List[Tuple[Any, ...]] = [(20, ), (2.7, ), (False, ), ] scripted_pdt_model = torch.jit.script(pdt_model, example_inputs={pdt_model: inp}) self.assertEqual(scripted_pdt_model(50), pdt_model(50)) self.assertEqual(scripted_pdt_model(1.8), pdt_model(1.8)) self.assertTrue(scripted_pdt_model(True), pdt_model(True)) def test_nested_nn_module_class(self): class NestedPDTInner(torch.nn.Module): def __init__(self): super().__init__() def forward(self, x): if isinstance(x, int): return x * 10 return x class NestedModulePDTWrapper(torch.nn.Module): def __init__(self, inner): super().__init__() self.inner = inner def forward(self, x): return self.inner(x) make_global(NestedPDTInner, NestedModulePDTWrapper) inner_pdt_model = NestedPDTInner() wrapped_pdt_model = NestedModulePDTWrapper(inner_pdt_model) inp: List[Tuple[Any, ...]] = [(20, ), (False, )] scripted_pdt_model = torch.jit.script(wrapped_pdt_model, example_inputs={wrapped_pdt_model: inp}) self.assertEqual(scripted_pdt_model(30), wrapped_pdt_model(30)) self.assertEqual(scripted_pdt_model(1.9), wrapped_pdt_model(1.9)) self.assertTrue(scripted_pdt_model(True), wrapped_pdt_model(True)) def test_nested_nn_module_class_with_args(self): class NestedModulePDTInner(torch.nn.Module): def __init__(self): super().__init__() def forward(self, x, y): if isinstance(x, int): return x * 10 + y return x class NestedModulePDTOuter(torch.nn.Module): def __init__(self, inner): super().__init__() self.inner = inner def forward(self, x): return self.inner(x, 20) make_global(NestedModulePDTInner, NestedModulePDTOuter) inner_pdt_model = NestedModulePDTInner() outer_pdt_model = NestedModulePDTOuter(inner_pdt_model) inner_input: List[Tuple[Any, ...]] = [(10, 10), (1.9, 20), ] outer_input: List[Tuple[Any, ...]] = [(20, ), (False, )] scripted_pdt_model = torch.jit.script(outer_pdt_model, example_inputs={inner_pdt_model: inner_input, outer_pdt_model: outer_input, }) self.assertEqual(scripted_pdt_model(30), outer_pdt_model(30)) self.assertEqual(scripted_pdt_model(1.9), outer_pdt_model(1.9)) self.assertTrue(scripted_pdt_model(True), outer_pdt_model(True)) def test_nested_function_in_forward(self): class NestedFunctionInForward(torch.nn.Module): def __init__(self): super().__init__() def forward(self, x): return self.fun(x) + 10 def fun(self, x): if isinstance(x, bool): return 0 elif isinstance(x, int): return x + 1 return 0 make_global(NestedFunctionInForward) pdt_model = NestedFunctionInForward() inp: List[Tuple[Any, ...]] = [(-1, ), (False, )] scripted_pdt_model = torch.jit.script(pdt_model, example_inputs={pdt_model: inp}) self.assertEqual(scripted_pdt_model(30), pdt_model(30)) self.assertEqual(scripted_pdt_model(True), pdt_model(True)) def test_nn_module_with_export_function(self): class TestModelWithExport(torch.nn.Module): def __init__(self): super().__init__() @torch.jit.export def fn(self, x, y) -> Any: assert not (isinstance(x, bool) and isinstance(y, bool)) if isinstance(x, int) and isinstance(y, int): return x + y elif isinstance(x, float) and isinstance(y, float): return x - y else: return -1 make_global(TestModelWithExport) pdt_model = TestModelWithExport() inp: List[Tuple[Any, ...]] = [(20, 10, ), (2.7, 8.9, ), ] scripted_pdt_model = torch.jit.script(pdt_model, example_inputs={pdt_model.fn: inp}) self.assertEqual(scripted_pdt_model.fn(10, 90), pdt_model.fn(10, 90)) self.assertEqual(scripted_pdt_model.fn(1.8, 2.2), pdt_model.fn(1.8, 2.2)) self.assertTrue(scripted_pdt_model.fn(torch.ones(1), 2), pdt_model.fn(torch.ones(1), 2)) def test_class_methods(self): class PDTModel: def test_sum(self, a): return sum(a) make_global(PDTModel) pdt_model = PDTModel() inp: List[Tuple[Any, ...]] = [([10, 20, ], ), ] scripted_pdt_model = torch.jit.script(PDTModel, example_inputs={pdt_model.test_sum: inp}) script_model = scripted_pdt_model() self.assertEqual(script_model.test_sum([10, 20, 30, ], ), pdt_model.test_sum([10, 20, 30, ], )) def test_class_with_multiple_methods(self): class PDTModelWithManyMethods: def test_list_to_dict(self, a): new_dictionary: Dict[float, bool] = {} for element in a: new_dictionary[element] = True return new_dictionary def test_substring(self, a, b): return b in a make_global(PDTModelWithManyMethods) pdt_model = PDTModelWithManyMethods() list_inp: List[Tuple[Any, ...]] = [([1.2, 2.3, ], ), ] str_inp: List[Tuple[Any, ...]] = [("abc", "b", ), ] scripted_pdt_model = torch.jit.script(PDTModelWithManyMethods, example_inputs={pdt_model.test_list_to_dict: list_inp, pdt_model.test_substring: str_inp}) script_model = scripted_pdt_model() self.assertEqual(script_model.test_list_to_dict([1.1, 2.2, 3.3, ], ), pdt_model.test_list_to_dict([1.1, 2.2, 3.3, ], )) self.assertEqual(script_model.test_substring("helloworld", "world", ), pdt_model.test_substring("helloworld", "world", )) self.assertEqual(script_model.test_substring("helloworld", "def", ), pdt_model.test_substring("helloworld", "def", )) def test_multiple_class_with_same_method(self): class PDTModelOne: def test_find(self, a, b): return b in a.keys() class PDTModelTwo: def test_find(self, a, b): return b in a make_global(PDTModelOne, PDTModelTwo) pdt_model_one = PDTModelOne() pdt_model_two = PDTModelTwo() dict_inp: List[Tuple[Any, ...]] = [({1.2: True, 2.3: False, }, 1.2), ] list_inp: List[Tuple[Any, ...]] = [(["abc", "b", ], "c"), ] scripted_pdt_model_one = torch.jit.script(PDTModelOne, example_inputs={pdt_model_one.test_find: dict_inp}) scripted_pdt_model_two = torch.jit.script(PDTModelTwo, example_inputs={pdt_model_two.test_find: list_inp}) script_model_one, script_model_two = scripted_pdt_model_one(), scripted_pdt_model_two() self.assertEqual(script_model_one.test_find({1.1: True, 2.2: True, 3.3: False, }, 4.4), pdt_model_one.test_find({1.1: True, 2.2: True, 3.3: False, }, 4.4)) self.assertEqual(script_model_two.test_find(["hello", "world", ], "world"), pdt_model_two.test_find(["hello", "world", ], "world")) def test_pdt(self): def test_sum(a, b): return a + b make_global(test_sum) scripted_fn_add = torch.jit.script(test_sum, example_inputs=[(3, 4)]) self.assertEqual(scripted_fn_add(10, 2), test_sum(10, 2)) def test_sub(a, b): return a - b make_global(test_sub) scripted_fn_sub = torch.jit.script(test_sub, example_inputs=[(3.9, 4.10)]) self.assertEqual(scripted_fn_sub(6.5, 2.9), test_sub(6.5, 2.9)) def test_mul(a, b): return a * b make_global(test_mul) scripted_fn_mul = torch.jit.script(test_mul, example_inputs=[(-10, 9)]) self.assertEqual(scripted_fn_mul(-1, 3), test_mul(-1, 3)) def test_args_complex(real, img): return torch.complex(real, img) make_global(test_args_complex) scripted_fn_complex = torch.jit.script(test_args_complex, example_inputs=[(torch.rand(3, 4), torch.rand(3, 4))]) arg1, arg2 = torch.rand(3, 4), torch.rand(3, 4) self.assertEqual(scripted_fn_complex(arg1, arg2), test_args_complex(arg1, arg2)) def test_bool(a): if a: return -1 else: return 0 make_global(test_bool) scripted_fn_bool = torch.jit.script(test_bool, example_inputs=[(True,)]) self.assertEqual(scripted_fn_bool(True), test_bool(True)) def test_str(a): if a == "": return False else: return True make_global(test_str) scripted_fn_str = torch.jit.script(test_str, example_inputs=[("",)]) self.assertEqual(scripted_fn_str("abc"), test_str("abc")) def test_pdt_list_and_tuple(self): def test_list_and_tuple(a): return sum(a) make_global(test_list_and_tuple) scripted_fn_float_list_input = torch.jit.script(test_list_and_tuple, example_inputs=[([4.9, 8.9],)]) self.assertEqual(scripted_fn_float_list_input([11.9, 7.6]), test_list_and_tuple([11.9, 7.6])) scripted_fn_bool_list_input = torch.jit.script(test_list_and_tuple, example_inputs=[([True, False, True],)]) self.assertEqual(scripted_fn_bool_list_input([True, True, True]), test_list_and_tuple([True, True, True])) scripted_fn_int_list_input = torch.jit.script(test_list_and_tuple, example_inputs=[([3, 4, 5], )]) self.assertEqual(scripted_fn_int_list_input([1, 2, 3]), test_list_and_tuple([1, 2, 3])) scripted_fn_float_tuple_input = torch.jit.script(test_list_and_tuple, example_inputs=[((4.9, 8.9),)]) self.assertEqual(scripted_fn_float_tuple_input((11.9, 7.6)), test_list_and_tuple((11.9, 7.6))) scripted_fn_bool_tuple_input = torch.jit.script(test_list_and_tuple, example_inputs=[((True, False, True),)]) self.assertEqual(scripted_fn_bool_tuple_input((True, True, True)), test_list_and_tuple((True, True, True))) scripted_fn_int_tuple_input = torch.jit.script(test_list_and_tuple, example_inputs=[((3, 4, 5), )]) self.assertEqual(scripted_fn_int_tuple_input((1, 2, 3)), test_list_and_tuple((1, 2, 3))) def test_nested_list_and_tuple(self): def test_nested_list(inp): return [sum(v) for v in inp] def test_nested_tuple(inp): ans = 0.0 for tup in inp: for val in tup: if val > 0: ans *= val return ans make_global(test_nested_list, test_nested_tuple) list_inp = [[1, 2, 3, ], [5, 6, 7, ]] scripted_fn = torch.jit.script(test_nested_list, example_inputs=[(list_inp, ), ]) inp = [[0, 4, 7, ], [8, 11, ], [6, -1, -20, ]] self.assertEqual(scripted_fn(inp, ), test_nested_list(inp, )) list_inp = ([1, 2, 3, ], [5, 6, 7, ]) scripted_fn = torch.jit.script(test_nested_list, example_inputs=[(list_inp, ), ]) inp = ([0, 4, 7, ], [8, 11, ], [6, -1, -20, ]) self.assertEqual(scripted_fn(inp, ), test_nested_list(inp, )) tup_inp = [(1.0, 2.6, 3.7, ), (5.7, 6.1, 1.7, )] scripted_fn = torch.jit.script(test_nested_tuple, example_inputs=[(tup_inp, ), ]) inp = [(1.0, 4.1, 7.4, ), (4.8, 1.1, -1.2, ), (6.3, -1.3, -2.0, )] self.assertEqual(scripted_fn(inp, ), test_nested_tuple(inp, )) tup_inp = ((True, False, True, ), (False, False, False, )) scripted_fn = torch.jit.script(test_nested_tuple, example_inputs=[(tup_inp, ), ]) inp = ((True, True, True, ), (False, False, True, )) self.assertEqual(scripted_fn(inp, ), test_nested_tuple(inp, )) def test_pdt_dict(self): def test_dict(a): return a['foo'] def test_dict_int_list(a): return a[1] make_global(test_dict, test_dict_int_list) str_bool_inp = {'foo' : True, 'bar': False} scripted_fn = torch.jit.script(test_dict, example_inputs=[(str_bool_inp,)]) self.assertEqual(scripted_fn({'foo' : False, 'bar': True}, ), test_dict({'foo' : False, 'bar': True}, )) str_list_inp = {0 : [True, False], 1: [False, True]} scripted_fn = torch.jit.script(test_dict_int_list, example_inputs=[(str_list_inp,)]) self.assertEqual(scripted_fn({0 : [False, False], 1: [True, True]}, ), test_dict_int_list({0 : [False, False], 1: [True, True]}, )) def test_any(self): def test_multiple_types(a): assert not isinstance(a, bool) return a def test_multiple_type_refinement(a): if isinstance(a, bool): return 1 elif isinstance(a, int): return 1 + a elif isinstance(a, float): return 1 + int(a) else: return -1 make_global(test_multiple_types, test_multiple_type_refinement) scripted_fn = torch.jit.script(test_multiple_types, example_inputs=[(1,), ("abc", ), (8.9,), ([3, 4, 5], )]) self.assertEqual(scripted_fn(10), test_multiple_types(10)) self.assertEqual(scripted_fn("def"), test_multiple_types("def")) self.assertEqual(scripted_fn(7.89999), test_multiple_types(7.89999)) self.assertEqual(scripted_fn([10, 11, 14]), test_multiple_types([10, 11, 14])) scripted_fn = torch.jit.script(test_multiple_type_refinement, example_inputs=[(1,), ("abc", ), (8.9,), ([3, 4, 5],), (True, ), ({"a": True}, ), ]) self.assertEqual(scripted_fn(10), test_multiple_type_refinement(10)) self.assertEqual(scripted_fn("def"), test_multiple_type_refinement("def")) self.assertEqual(scripted_fn(7.89999), test_multiple_type_refinement(7.89999)) self.assertEqual(scripted_fn([10, 11, 14]), test_multiple_type_refinement([10, 11, 14])) self.assertEqual(scripted_fn(False), test_multiple_type_refinement(False)) self.assertEqual(scripted_fn({"abc" : True, "def": False}), test_multiple_type_refinement({"abc" : True, "def": False})) def test_class_as_profiled_types(self): class UserDefinedClass: def fn(self, b) -> Any: assert b is not None if isinstance(b, int): return b if b > 0 else -1 elif isinstance(b, float): return b if b > 0.0 else -1.0 return 0 def test_model(a, m): assert not isinstance(a, bool) return m.fn(a) make_global(UserDefinedClass, test_model) user_class = UserDefinedClass() scripted_fn = torch.jit.script(test_model, example_inputs=[(10, user_class, ), (10.9, user_class, ), ]) self.assertEqual(scripted_fn(100, user_class, ), test_model(100, user_class)) self.assertEqual(scripted_fn(1.9, user_class, ), test_model(1.9, user_class)) def test_class_with_args_as_profiled_types(self): class ClassWithArgs: def __init__(self, a: bool): self.a = a def fn(self, b): if self.a: return b else: return -1 def test_model_with_args(a, m): assert not isinstance(a, bool) return m.fn(a) make_global(ClassWithArgs, test_model_with_args) user_class = ClassWithArgs(False) scripted_fn = torch.jit.script(test_model_with_args, example_inputs=[(10, user_class, ), (10.9, user_class, ), ]) self.assertEqual(scripted_fn(100, ClassWithArgs(True), ), test_model_with_args(100, ClassWithArgs(True))) def test_nn_parameter_as_arg(self): class TestNNParameter(torch.nn.Module): def __init__(self): super().__init__() self.inp = torch.nn.Parameter(torch.ones(2, 3)) def add_nn_parameter_with_int(self, x, y): return torch.add(x, y) def forward(self, y): return self.add_nn_parameter_with_int(self.inp, y) make_global(TestNNParameter) pdt_model = TestNNParameter() scripted_fn = torch.jit.script(pdt_model, example_inputs={pdt_model: [(10, ), ], }) self.assertEqual(scripted_fn(20), pdt_model(20)) def test_fx_tracing_with_typing(self): class FXModelOutput(NamedTuple): result: List[int] class FXModel(torch.nn.Module): def forward(self, a) -> FXModelOutput: result = FXModelOutput(result=a) return result make_global(FXModel, FXModelOutput) pdt_model = FXModel() scripted_fn = torch.jit.script(pdt_model, example_inputs={pdt_model: [([10, 20, ], ), ], }) self.assertEqual(scripted_fn([20]), pdt_model([20])) def test_nonetype_as_optional_of_type(self): def test_none(a) -> Any: if a is None: return 0 else: return a + torch.ones(1) make_global(test_none) scripted_fn = torch.jit.script(test_none, example_inputs=[(None, ), (10.6, )]) self.assertEqual(scripted_fn(30.9, ), test_none(30.9, )) scripted_fn = torch.jit.script(test_none, example_inputs=[(None, ), (10, )]) self.assertEqual(scripted_fn(2, ), test_none(2, )) scripted_fn = torch.jit.script(test_none, example_inputs=[(None, ), (torch.Tensor(1), )]) self.assertEqual(scripted_fn(torch.ones(1), ), test_none(torch.ones(1), ))