| # Owner(s): ["oncall: jit"] |
| |
| import os |
| import sys |
| from typing import Any, Dict, List, NamedTuple, Optional, Tuple # noqa: F401 |
| |
| import torch |
| from torch.jit._monkeytype_config import _IS_MONKEYTYPE_INSTALLED |
| from torch.testing._internal.common_utils import NoTest |
| from torch.testing._internal.jit_utils import JitTestCase, make_global |
| |
| |
| # 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 = NoTest # 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 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 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 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 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): |
| @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) -> None: |
| 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), |
| ), |
| ) |
