| # Owner(s): ["oncall: jit"] |
| |
| import io |
| import os |
| import sys |
| |
| import torch |
| from torch.testing import FileCheck |
| from enum import Enum |
| from textwrap import dedent |
| from typing import Dict, List, Optional, Tuple, Union |
| |
| # 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) |
| from torch.testing._internal.jit_utils import JitTestCase, make_global |
| |
| 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 TestUnion(JitTestCase): |
| """ |
| This class tests the functionality of `Union`. |
| |
| Note: It's important to be able to refine the type of a `Union` to |
| one of its internal types. Currently, there are differences in the |
| way Python expects `isinstance` checks and the way TorchScript |
| expects `isinstance` checks. This means that we can't use |
| `checkScript` in our test cases because either the eager mode or the |
| script mode wouldn't run! So, some test cases have separate but |
| equivalent functions to emulate `checkScript`. |
| """ |
| |
| def test_check_union_annotation(self): |
| def test_func(a: Union[int, float], b: Optional[int]): |
| return 0 |
| |
| scripted_func = torch.jit.script(test_func) |
| graph_rep = str(scripted_func.graph) |
| code_rep = str(scripted_func.code) |
| # TS graph IR for Union should be annotated as Union() |
| FileCheck().check("Union(").check("int?").run(graph_rep) |
| # Serialized code for Union should be annotated as Union[] |
| FileCheck().check("Union[").check("Optional[int]").run(code_rep) |
| self.checkScript(test_func, (5, 6)) |
| # this shouldn't error out |
| torch._C.parse_ir(str(scripted_func.graph)) |
| |
| def test_union_with_scalar_values(self): |
| def fn(x: Union[int, float]) -> str: |
| return "foo" |
| |
| self.checkScript(fn, (1,)) |
| self.checkScript(fn, (1.0,)) |
| |
| scripted = torch.jit.script(fn) |
| |
| with self.assertRaisesRegex(RuntimeError, "Expected a member of" |
| r" Union\[float, int\] but " |
| "instead found type str"): |
| scripted("1") |
| |
| def test_union_with_collections(self): |
| def fn(x: Union[Dict[str, int], List[int]]) -> str: |
| return "foo" |
| |
| self.checkScript(fn, ({"foo": 1, "bar": 2, "baz": 3},)) |
| self.checkScript(fn, ([1, 2, 3],)) |
| |
| scripted = torch.jit.script(fn) |
| |
| with self.assertRaisesRegex(RuntimeError, "Expected a member of" |
| r" Union\[List\[int\], Dict\[str, " |
| r"int\]\] but instead found type " |
| r"Dict\[str, str\]"): |
| scripted({"foo": "bar", "baz": "qux"}) |
| |
| with self.assertRaisesRegex(RuntimeError, "Expected a member of" |
| r" Union\[List\[int\], Dict\[str, " |
| r"int\]\] but instead found type " |
| r"List\[str\]"): |
| scripted(["foo", "bar", "baz"]) |
| |
| with self.assertRaisesRegex(RuntimeError, "Expected a member of" |
| r" Union\[List\[int\], Dict\[str, " |
| r"int\]\] but instead found type " |
| "str"): |
| scripted("1") |
| |
| def test_union_with_enum(self): |
| class Color(Enum): |
| RED = 1 |
| GREEN = 2 |
| |
| make_global(Color) |
| |
| def fn(x: Union[str, Color]) -> str: |
| return "foo" |
| |
| self.checkScript(fn, (Color.RED,)) |
| self.checkScript(fn, ("red",)) |
| |
| scripted = torch.jit.script(fn) |
| |
| with self.assertRaisesRegex(RuntimeError, "Expected a member of" |
| r" Union\[__torch__.jit.test_union." |
| r"Color, str\] but instead found " |
| "type int"): |
| scripted(1) |
| |
| def test_union_in_class_constructor(self): |
| |
| @torch.jit.script # noqa: B903 |
| class A(object): # noqa: B903 |
| def __init__(self, x: Union[int, str]) -> None: |
| self.x = x |
| |
| def fn(x: Union[str, int]) -> A: |
| return A(x) |
| |
| self.assertEqual(fn("foo").x, "foo") |
| self.assertEqual(fn(1).x, 1) |
| |
| scripted = torch.jit.script(fn) |
| |
| with self.assertRaisesRegex(RuntimeError, "Expected a member of" |
| r" Union\[int, str\] but instead " |
| r"found type List\[str\]"): |
| scripted(["foo", "bar", "baz"]) |
| |
| def test_union_return_type(self): |
| def fn(x: int) -> Union[int, str]: |
| return "foo" |
| |
| self.checkScript(fn, (1,)) |
| |
| def test_union_as_annotation(self): |
| def fn() -> Union[int, str]: |
| x: Union[int, str] = "foo" |
| return x |
| |
| self.checkScript(fn, ()) |
| |
| def test_union_as_annotation_in_typed_container(self): |
| def fn() -> None: |
| l: List[Union[int, str]] = [] |
| u1: Union[int, str] = "foo" |
| u2: Union[int, str] = 1 |
| l.append(u1) |
| l.append(u2) |
| |
| self.checkScript(fn, ()) |
| |
| def test_union_as_annotation_py2(self): |
| def fn(): |
| # type: () -> Union[int, str] |
| x: Union[int, str] = "foo" |
| return x |
| |
| self.checkScript(fn, ()) |
| |
| def test_union_as_internal_tuple_type(self): |
| def fn(): |
| t: Tuple[Union[int, str], Union[int, str]] = (1, "foo") |
| return t |
| |
| self.checkScript(fn, ()) |
| |
| def test_union_variable_can_be_reassigned(self): |
| @torch.jit.script |
| def aux1(i: int): |
| return int(i ** 2) |
| |
| @torch.jit.script |
| def aux2(s: str): |
| return s + s |
| |
| def fn() -> Union[int, str]: |
| x: Union[int, str] = "foo" |
| i: int = 1 |
| x = i |
| y: int = aux1(x) |
| z: str = aux2(str(y)) |
| x = z |
| return x |
| |
| self.checkScript(fn, ()) |
| |
| def test_union_does_not_replace_existing_annotated_type(self): |
| def fn(): |
| x: List[int] = [1, 2, 3] |
| x.append("foo") |
| return x |
| |
| with self.assertRaisesRegex(RuntimeError, "Could not match type str"): |
| scripted = torch.jit.script(fn) |
| scripted() |
| |
| def test_union_does_not_replace_existing_annotated_type_union(self): |
| def fn(): |
| x: List[Union[int, str]] = [1, "foo", 3] |
| x.append(2.0) |
| return x |
| |
| with self.assertRaisesRegex(RuntimeError, "Could not match type float"): |
| scripted = torch.jit.script(fn) |
| scripted() |
| |
| def test_union_does_not_replace_existing_annotated_type_empty_container(self): |
| def fn(): |
| x: List[int] = [] |
| x.append("foo") |
| return x |
| |
| with self.assertRaisesRegex(RuntimeError, "Could not match type str"): |
| scripted = torch.jit.script(fn) |
| scripted() |
| |
| def test_unions_of_unions_are_flattened(self): |
| @torch.jit.script |
| def fn(x: Union[Union[int, str], float]) -> str: |
| return "foo" |
| |
| s = fn.graph |
| |
| FileCheck().check("x : Union(float, int, str)") \ |
| .run(s) |
| |
| def test_unions_of_a_single_argument_vanish(self): |
| @torch.jit.script |
| def fn(x: Union[int]) -> str: |
| return "foo" |
| |
| s = fn.graph |
| |
| FileCheck().check("x : int") \ |
| .run(s) |
| |
| def test_union_redundant_arguments_are_skipped(self): |
| @torch.jit.script |
| def fn(x: Union[int, str, int]) -> str: |
| return "foo" |
| |
| s = fn.graph |
| |
| FileCheck().check("x : Union(int, str)") \ |
| .run(s) |
| |
| def test_union_redundant_arguments_are_skipped_optional(self): |
| @torch.jit.script |
| def fn(x: Union[int, Optional[float], Optional[int]]) -> str: |
| return "foo" |
| |
| s = fn.graph |
| |
| FileCheck().check("x : Union(float, int, NoneType)") \ |
| .run(s) |
| |
| def test_union_redundant_arguments_are_skipped_subtyping(self): |
| @torch.jit.script |
| def fn(x: Union[str, Tuple[Optional[int], int], Tuple[int, int]]) -> str: |
| return "foo" |
| |
| s = fn.graph |
| |
| FileCheck().check("x : Union((int?, int), str)") \ |
| .run(s) |
| |
| def test_union_redundant_arguments_are_skipped_container(self): |
| @torch.jit.script |
| def fn(x: Union[List[str], List[float], List[str]]) -> str: |
| return "foo" |
| |
| s = fn.graph |
| |
| FileCheck().check("x : Union(float[], str[])") \ |
| .run(s) |
| |
| def test_union_argument_order_is_ignored(self): |
| @torch.jit.script |
| def fn1(x: Union[int, str]) -> str: |
| return "foo" |
| |
| @torch.jit.script |
| def fn2(x: Union[str, int]) -> str: |
| return "foo" |
| |
| for s in (fn1.graph, fn2.graph): |
| FileCheck().check("x : Union(int, str)") \ |
| .run(s) |
| |
| def test_union_argument_order_is_ignored_container(self): |
| @torch.jit.script |
| def fn1(x: Union[List[str], List[int]]) -> str: |
| return "foo" |
| |
| @torch.jit.script |
| def fn2(x: Union[List[int], List[str]]) -> str: |
| return "foo" |
| |
| for s in (fn1.graph, fn2.graph): |
| FileCheck().check("x : Union(int[], str[])") \ |
| .run(s) |
| |
| def test_union_T_None_is_equivalent_to_optional_T(self): |
| @torch.jit.script |
| def inner(x: Union[int, None]) -> int: |
| if x is not None: |
| return x |
| else: |
| return 5 |
| |
| @torch.jit.script |
| def fn1() -> int: |
| a: Optional[int] = 5 |
| b: Optional[int] = None |
| a_ = inner(a) |
| b_ = inner(b) |
| return a_ + b_ |
| |
| self.assertEqual(fn1(), 10) |
| |
| @torch.jit.script |
| def inner2(x: Optional[int]) -> int: |
| if x is not None: |
| return x |
| else: |
| return 5 |
| |
| @torch.jit.script |
| def fn2() -> int: |
| a: Union[int, None] = 5 |
| b: Union[int, None] = None |
| a_ = inner(a) |
| b_ = inner(b) |
| return a_ + b_ |
| |
| self.assertEqual(fn2(), 10) |
| |
| def test_union_optional_of_union_is_flattened(self): |
| @torch.jit.script |
| def fn(flag: int) -> Union[str, int, None]: |
| y: Union[int, str, None] = "foo" |
| if flag == 0: |
| x: Optional[Union[int, str]] = y |
| elif flag == 1: |
| x: Optional[Union[int, str]] = 1 |
| else: |
| x: Optional[Union[int, str]] = None |
| return x |
| |
| # Can't use `checkScript` because it will flag the fact that |
| # the original code has `Optional[Union[int, str]]` but the |
| # saved/loaded code has `Union[int, NoneType, str]` (even |
| # though this is exactly what we want) |
| self.assertEqual(fn(0), "foo") |
| self.assertEqual(fn(1), 1) |
| self.assertEqual(fn(2), None) |
| |
| buffer = io.BytesIO() |
| torch.jit.save(fn, buffer) |
| buffer = io.BytesIO(buffer.getvalue()) |
| l = torch.jit.load(buffer) |
| |
| s = l.code |
| |
| FileCheck().check("Union[int, NoneType, str]") \ |
| .check("Union[int, NoneType, str]") \ |
| .run(s) |
| |
| def test_union_subclasses_larger_union(self): |
| def fn() -> Union[int, str, torch.Tensor]: |
| x: Union[int, str] = "foo" |
| return x |
| |
| self.checkScript(fn, ()) |
| |
| # TODO: We would like to eventually support this. The issue is being |
| # tracked at https://github.com/pytorch/pytorch/issues/58167 |
| def test_union_as_dict_key(self): |
| def fn(): |
| x: Dict[Union[int, str], str] = {} |
| x["foo"] = "bar" |
| x[1] = 2 |
| return x[1] |
| |
| with self.assertRaisesRegex(RuntimeError, "only int, float, " |
| "complex, Tensor, device and string keys " |
| "are supported"): |
| torch.jit.script(fn) |
| |
| def test_union_as_dict_value(self): |
| def fn(): |
| x: Dict[str, Union[int, str]] = {} |
| x["foo"] = "bar" |
| x["baz"] = 2 |
| return x["baz"] |
| |
| self.checkScript(fn, ()) |
| |
| def test_union_module_with_union_instance_variable(self): |
| class M(torch.nn.Module): |
| |
| x: Union[int, str] |
| |
| def __init__(self, x: Union[int, str]): |
| super().__init__() |
| self.x: Union[int, str] = x |
| |
| def forward(self, y: Union[int, str]): |
| self.x = y |
| return self.x |
| |
| self.checkModule(M(2,), (1,)) |
| self.checkModule(M("bar"), ("foo",)) |
| |
| def test_union_module_with_union_class_variable(self): |
| class M(torch.nn.Module): |
| x: Union[int, str] = "foo" |
| |
| def __init__(self, y: int): |
| super().__init__() |
| x = y |
| |
| def forward(self, z: str): |
| x = z |
| return x |
| |
| self.checkModule(M(1), ("foo",)) |
| |
| def test_union_type_refinement(self): |
| def fn(x: Union[int, str]) -> str: |
| if isinstance(x, str): |
| z = x + "bar" |
| return x |
| else: |
| return "baz" |
| |
| self.checkScript(fn, ("foo",)) |
| self.checkScript(fn, (1,)) |
| |
| def test_union_type_refinement_union_rhs(self): |
| def fn(x: int) -> str: |
| if torch.jit.isinstance(x, Union[int, str]): |
| return "bar" |
| else: |
| return "baz" |
| |
| self.checkScript(fn, (1,)) |
| |
| def test_union_type_refinement_tuple_rhs(self): |
| def fn(x: Union[int, float, List[str]]) -> str: |
| if isinstance(x, (int, float)): |
| if isinstance(x, int): |
| return str(x) |
| else: |
| return "foo" |
| else: |
| if len(x): |
| return x[0] |
| else: |
| return "bar" |
| |
| self.checkScript(fn, (1,)) |
| self.checkScript(fn, (1.0,)) |
| self.checkScript(fn, (["a", "b", "c"],)) |
| |
| def test_union_type_refinement_tuple_rhs_noncontained_type(self): |
| def fn(x: Union[int, List[str]]) -> str: |
| if isinstance(x, (int, float)): |
| y = x + x |
| return str(y) |
| else: |
| if len(x): |
| return x[0] |
| else: |
| return "bar" |
| |
| self.checkScript(fn, (1,)) |
| self.checkScript(fn, (["a", "b", "c"],)) |
| |
| def test_union_type_refinement_tuple_rhs_union(self): |
| @torch.jit.script |
| def fn(x: int) -> str: |
| if torch.jit.isinstance(x, (Union[int, str], float)): |
| y = x + x |
| return str(y) |
| else: |
| return "foo" |
| |
| # TODO: There's currently an unrelated bug in |
| # `torch.jit.isinstance` that makes it fail for tuple literals. |
| # Posted here: https://github.com/pytorch/pytorch/issues/60095 |
| # Change `assertEqual` to `checkScript` when the bug is fixed |
| self.assertEqual(fn(1), "2") |
| |
| def test_union_type_refinement_statically_false(self): |
| @torch.jit.script |
| def fn(x: int) -> str: |
| if torch.jit.isinstance(x, (Union[str, float], List[str], str)): |
| z = x + "foo" |
| return z |
| else: |
| return "bar" |
| |
| s = fn.graph |
| |
| # Check that we don't have any branching statements |
| FileCheck().check_not("block0()") \ |
| .check_not("block1()") \ |
| .run(s) |
| |
| def test_union_type_refinement_statically_true(self): |
| @torch.jit.script |
| def fn(x: Union[List[int], int]) -> Union[List[int], int]: |
| if not torch.jit.isinstance(x, (int, List[int])): |
| return x |
| else: |
| l = [1, 2, 3] |
| y: Union[List[int], int] = l |
| return y |
| |
| s = fn.graph |
| |
| # Check that we don't have any branching statements |
| FileCheck().check_not("block0()") \ |
| .check_not("block1()") \ |
| .run(s) |
| |
| def test_union_type_refinement_partial_static_refinement_tuple_rhs(self): |
| def fn(x: Union[List[int], int]) -> int: |
| if torch.jit.isinstance(x, (int, float, str)): |
| # We should know that `x` is an `int` here |
| z = x + 1 |
| return z |
| else: |
| return 100 |
| |
| self.checkScript(fn, ([1, 2, 3],)) |
| self.checkScript(fn, (1,)) |
| |
| def test_union_type_refinement_partial_static_refinement_union_rhs(self): |
| def fn(x: Union[List[int], int]) -> int: |
| if torch.jit.isinstance(x, Union[int, float, str]): |
| # We should know that `x` is an `int` here |
| z = x + 1 |
| return z |
| else: |
| return 100 |
| |
| self.checkScript(fn, ([1, 2, 3],)) |
| self.checkScript(fn, (1,)) |
| |
| def test_union_type_refinement_internal_declaration(self): |
| def fn(flag: bool) -> str: |
| x: Union[int, str, None] = None |
| if (flag): |
| y = "foo" |
| else: |
| y = 1 |
| if isinstance(x, str): |
| return x |
| else: |
| return "bar" |
| |
| self.checkScript(fn, (True,)) |
| self.checkScript(fn, (False,)) |
| |
| def test_union_branching_with_union_return_and_homogenous_types(self): |
| def fn(x: int) -> Union[int, str]: |
| if x % 2: |
| return "foo" |
| else: |
| return "bar" |
| |
| self.checkScript(fn, (1,)) |
| self.checkScript(fn, (8,)) |
| |
| def test_union_branching_does_not_autoinfer_undeclared_union(self): |
| def fn(x: int) -> str: |
| if x % 2: |
| y = "foo" |
| else: |
| y = x |
| if isinstance(y, str): |
| return y |
| else: |
| return "bar" |
| |
| with self.assertRaisesRegex(RuntimeError, "y is set to type str" |
| " in the true branch and type int " |
| "in the false branch"): |
| torch.jit.script(fn) |
| |
| def test_union_branching_does_not_widen_existing_inferred_type(self): |
| def fn(x: int) -> str: |
| y = "foo" |
| if x % 2: |
| y = "bar" |
| else: |
| y = x |
| if isinstance(y, str): |
| return y |
| else: |
| return "baz" |
| |
| with self.assertRaisesRegex(RuntimeError, "previously had type " |
| "str but is now being assigned to a" |
| " value of type int"): |
| torch.jit.script(fn) |
| |
| def test_union_schema_matching_on_internal_type(self): |
| def fn(x: Union[List[int], Dict[str, int]]) -> int: |
| if torch.jit.isinstance(x, List[int]): |
| return x[0] |
| else: |
| return list(x.values())[0] |
| |
| self.checkScript(fn, ([1, 2, 3],)) |
| self.checkScript(fn, ({"foo": 1, "bar": 2, "baz": 3},)) |
| |
| def test_union_subtractive_refinement(self): |
| def fn(x: Union[List[int], int]) -> int: |
| if not isinstance(x, int): |
| x.append(1) |
| return x[0] |
| else: |
| return x |
| |
| self.checkScript(fn, (1,)) |
| self.checkScript(fn, ([1, 2, 3],)) |
| |
| def test_union_subtractive_refinement_with_container(self): |
| def fn(x: Union[List[int], int]) -> int: |
| if not torch.jit.isinstance(x, List[int]): |
| return x |
| else: |
| x.append(1) |
| return x[0] |
| |
| self.checkScript(fn, (1,)) |
| self.checkScript(fn, ([1, 2, 3],)) |
| |
| def test_union_memory_aliasing(self): |
| def fn(): |
| x : List[torch.Tensor] = [] |
| z : List[Optional[List[torch.Tensor]]] = [] |
| z.append(x) |
| x_alias = z[0] |
| if torch.jit.isinstance(x_alias, List[torch.Tensor]): |
| x_alias.append(torch.tensor(3)) |
| return x |
| |
| self.checkScript(fn, ()) |
| |
| def test_union_serialization_preserves_type_annotations(self): |
| # This function will fail after being torch.jit.save'd and |
| # torch.jit.load'd if the type annotations aren't preserved |
| # for Union during serialization. We need the `Union[str, int]` |
| # annotation to make sure that `y` is typed as a Union instead |
| # of as a str in one branch and an int in the other |
| def fn(x: int) -> str: |
| if x % 2: |
| y: Union[str, int] = "bar" |
| else: |
| y: Union[str, int] = x |
| if isinstance(y, str): |
| return y |
| else: |
| return "baz" |
| |
| self.checkScript(fn, (1,)) |
| self.checkScript(fn, (8,)) |
| |
| def _assert_passes(self, template: str, ann: str, lhs: str): |
| code = template.format(ann=ann, lhs=lhs) |
| self.checkScript(code, (), name="fn") |
| |
| def _assert_raises(self, template: str, ann: str, lhs: str, msg: str): |
| code = template.format(ann=ann, lhs=lhs) |
| with self.assertRaisesRegex(RuntimeError, msg): |
| cu = torch.jit.CompilationUnit(code, _frames_up=1) |
| string_frontend = getattr(cu, "fn") # noqa: B009 |
| |
| def test_union_with_list_assignment(self): |
| template = dedent(''' |
| def fn(): |
| x: {ann} = {lhs} |
| if torch.jit.isinstance(x, List[torch.Tensor]): |
| x.append(torch.tensor(3)) |
| return x |
| ''') |
| |
| lhs = {"list_literal_empty" : "[]", |
| |
| "list_literal_of_tensor" : "[torch.arange(3), torch.arange(5)]", |
| |
| "list_literal_of_str" : "[\"foo\", \"bar\", \"baz\"]", |
| |
| "list_literal_of_mixed" : "[torch.arange(5), 1]", |
| |
| "list_comprehension_of_tensor" : |
| "[torch.add(x, 1) for x in [torch.arange(3), torch.arange(5)]]", |
| |
| "list_comprehension_of_str" : |
| "[x + \"!\" for x in [\"foo\", \"bar\", \"baz\"]]", |
| |
| "list_comprehension_of_mixed" : |
| "[torch.add(1, x) for x in [torch.arange(5), 1]]"} |
| |
| """ |
| Union[List[str], List[torch.Tensor]] |
| """ |
| self._assert_raises(template, |
| "Union[List[str], List[torch.Tensor]]", |
| lhs["list_literal_empty"], |
| "there are multiple possible List type " |
| "candidates in the Union annotation") |
| |
| self._assert_passes(template, |
| "Union[List[str], List[torch.Tensor]]", |
| lhs["list_literal_of_tensor"]) |
| |
| self._assert_passes(template, |
| "Union[List[str], List[torch.Tensor]]", |
| lhs["list_literal_of_str"]) |
| |
| self._assert_raises(template, |
| "Union[List[str], List[torch.Tensor]]", |
| lhs["list_literal_of_mixed"], |
| "none of those types match the types of the" |
| " given list elements") |
| |
| self._assert_passes(template, |
| "Union[List[str], List[torch.Tensor]]", |
| lhs["list_comprehension_of_tensor"]) |
| |
| self._assert_passes(template, |
| "Union[List[str], List[torch.Tensor]]", |
| lhs["list_comprehension_of_str"]) |
| |
| # TODO: Support mixed list comprehensions |
| self._assert_raises(template, |
| "Union[List[str], List[torch.Tensor]]", |
| lhs["list_comprehension_of_mixed"], |
| "Arguments for call are not valid") |
| |
| """ |
| Union[int, torch.Tensor] |
| """ |
| self._assert_raises(template, |
| "Union[int, torch.Tensor]", |
| lhs["list_literal_empty"], |
| "Expected an Union type annotation with an " |
| "inner List type") |
| |
| self._assert_raises(template, "Union[int, torch.Tensor]", |
| lhs["list_literal_of_tensor"], |
| "Expected an Union type annotation with an " |
| "inner List type") |
| |
| self._assert_raises(template, "Union[int, torch.Tensor]", |
| lhs["list_comprehension_of_tensor"], |
| "Expected an Union type annotation with an " |
| "inner List type") |
| |
| """ |
| Union[List[torch.Tensor], int] |
| """ |
| self._assert_passes(template, |
| "Union[List[torch.Tensor], int]", |
| lhs["list_literal_empty"]) |
| |
| self._assert_passes(template, |
| "Union[List[torch.Tensor], int]", |
| lhs["list_literal_of_tensor"]) |
| |
| self._assert_raises(template, "Union[List[torch.Tensor], int]", |
| lhs["list_literal_of_str"], |
| r"List type annotation `List\[Tensor\]` did " |
| "not match the types of the given list " |
| "elements") |
| |
| self._assert_raises(template, "Union[List[torch.Tensor], int]", |
| lhs["list_literal_of_mixed"], |
| r"List type annotation `List\[Tensor\]` did " |
| "not match the types of the given list " |
| "elements") |
| |
| self._assert_passes(template, |
| "Union[List[torch.Tensor], int]", |
| lhs["list_comprehension_of_tensor"]) |
| |
| self._assert_raises(template, |
| "Union[List[torch.Tensor], int]", |
| lhs["list_comprehension_of_str"], |
| r"List type annotation `List\[Tensor\]` did " |
| "not match the types of the given list " |
| "elements") |
| |
| # TODO(@ansley): Support mixed list comprehensions |
| self._assert_raises(template, |
| "Union[List[torch.Tensor], int]", |
| lhs["list_comprehension_of_mixed"], |
| "Arguments for call are not valid") |
| |
| def test_union_with_dict_assignment(self): |
| template = dedent(''' |
| def fn(): |
| x: {ann} = {lhs} |
| if torch.jit.isinstance(x, Dict[str, torch.Tensor]): |
| x["foo"] = torch.tensor(3) |
| return x |
| ''') |
| |
| lhs = {"dict_literal_empty" : "{}", |
| |
| "dict_literal_of_str_tensor" : |
| "{\"foo\" : torch.arange(3), \"bar\" : torch.arange(5)}", |
| |
| "dict_literal_of_str_int" : |
| "{\"foo\" : 1, \"bar\" : 2}", |
| |
| "dict_literal_of_mixed" : |
| "{\"foo\" : torch.arange(3), \"bar\" : 2}", |
| |
| "dict_comprehension_of_str_tensor" : |
| "{x : torch.add(y, 1) for x, y in \ |
| zip([\"foo\", \"bar\"], [torch.arange(3), torch.arange(5)])}", |
| |
| "dict_comprehension_of_str_int" : |
| "{x : torch.add(y, 1) for x, y in \ |
| zip([\"foo\", \"bar\"], [1, 2]}", |
| |
| "dict_comprehension_of_mixed" : |
| "{x : torch.add(y, 1) for x, y in \ |
| zip([\"foo\", \"bar\"], [torch.arange(3), 2])}", |
| |
| "dict_keyword" : |
| "dict(foo=torch.arange(3), baz=torch.arange(5))", |
| |
| "dict_keyword_with_iterable" : |
| "dict([(\"foo\", torch.arange(3)), (\"bar\", torch.arange(5))])", |
| |
| "dict_keyword_with_empty_iterable" : |
| "dict([])", |
| |
| "dict_keyword_with_internal_aggregate_function" : |
| "dict(zip([\"foo\", \"bar\"], [torch.arange(3), torch.arange(5)])", |
| |
| "dict_keyword_with_mapping" : |
| "dict({\"foo\" : torch.arange(3), \"bar\" : torch.arange(5)})", |
| |
| "dict_keyword_with_mapping_and_kwargs" : |
| "dict({\"foo\" : torch.arange(3), \"bar\" : torch.arange(5)}, baz=torch.arange(7))", |
| |
| } |
| |
| """ |
| Union[Dict[str, torch.Tensor], Dict[str, int]] |
| """ |
| self._assert_raises(template, |
| "Union[List[str], List[torch.Tensor]]", |
| lhs["dict_literal_empty"], |
| "Expected an Union type annotation with an " |
| "inner Dict type") |
| |
| self._assert_passes(template, |
| "Union[Dict[str, torch.Tensor], Dict[str, int]]", |
| lhs["dict_literal_of_str_tensor"]) |
| |
| self._assert_passes(template, |
| "Union[Dict[str, torch.Tensor], Dict[str, int]]", |
| lhs["dict_literal_of_str_int"]) |
| |
| self._assert_raises(template, "Union[Dict[str, torch.Tensor], Dict[str, int]]", |
| lhs["dict_literal_of_mixed"], |
| "none of those dict types can hold the " |
| "types of the given keys and values") |
| |
| # TODO: String frontend does not support tuple unpacking |
| # https://github.com/pytorch/pytorch/issues/64096 |
| # self._assert_passes(template, "Union[Dict[str, torch.Tensor], Dict[str, int]]", |
| # lhs["dict_comprehension_of_str_tensor"]) |
| |
| # self._assert_passes(template, "Union[Dict[str, torch.Tensor], Dict[str, int]]", |
| # lhs["dict_comprehension_of_str_int"]) |
| |
| # self._assert_raises(template, "Union[Dict[str, torch.Tensor], Dict[str, int]]", |
| # lhs["dict_comprehension_of_mixed"], |
| # "foobar") |
| |
| # self._assert_passes(template, |
| # "Union[Dict[str, torch.Tensor], Dict[str, int]]", |
| # lhs["dict_keyword_with_internal_aggregate_function"]) |
| |
| # TODO(@ansley): Follow-up project needed for full type |
| # inference with dict keyword (supported for dict comprehension |
| # and dict literal already; should not be a blocker for anyone) |
| self._assert_raises(template, |
| "Union[Dict[str, torch.Tensor], Dict[str, int]]", |
| lhs["dict_keyword"], |
| "full type inference is not yet supported") |
| |
| self._assert_raises(template, |
| "Union[Dict[str, torch.Tensor], Dict[str, int]]", |
| lhs["dict_keyword_with_iterable"], |
| "full type inference is not yet supported") |
| |
| self._assert_raises(template, |
| "Union[Dict[str, torch.Tensor], Dict[str, int]]", |
| lhs["dict_keyword_with_empty_iterable"], |
| "full type inference is not yet supported") |
| |
| self._assert_raises(template, |
| "Union[Dict[str, torch.Tensor], Dict[str, int]]", |
| lhs["dict_keyword_with_mapping"], |
| "full type inference is not yet supported") |
| |
| self._assert_raises(template, |
| "Union[Dict[str, torch.Tensor], Dict[str, int]]", |
| lhs["dict_keyword_with_mapping_and_kwargs"], |
| "full type inference is not yet supported") |
| |
| """ |
| Union[int, torch.Tensor] |
| """ |
| self._assert_raises(template, |
| "Union[int, torch.Tensor]", |
| lhs["dict_literal_empty"], |
| "Expected an Union type annotation with " |
| "an inner Dict type") |
| |
| self._assert_raises(template, |
| "Union[int, torch.Tensor]", |
| lhs["dict_literal_of_str_tensor"], |
| "Expected an Union type annotation with " |
| "an inner Dict type") |
| |
| # See above--string frontend does not support tuple unpacking |
| # self._assert_raises(template, "Union[int, torch.Tensor]", |
| # lhs["dict_comprehension_of_tensor"], |
| # "foobar") |
| |
| """ |
| Union[Dict[str, torch.Tensor], int] |
| """ |
| self._assert_passes(template, |
| "Union[Dict[str, torch.Tensor], int]", |
| lhs["dict_literal_empty"]) |
| |
| self._assert_passes(template, |
| "Union[Dict[str, torch.Tensor], int]", |
| lhs["dict_literal_of_str_tensor"]) |
| |
| self._assert_raises(template, |
| "Union[Dict[str, torch.Tensor], int]", |
| lhs["dict_literal_of_str_int"], |
| "Type annotation was inferred to be " |
| r"`Dict\[str, Tensor\]`, but the type of " |
| "values given by the dict literal is") |
| |
| self._assert_raises(template, |
| "Union[Dict[str, torch.Tensor], int]", |
| lhs["dict_literal_of_mixed"], |
| "Type annotation was inferred to be " |
| r"`Dict\[str, Tensor\]`, but the type of " |
| "values given by the dict literal is") |
| |
| self._assert_passes(template, |
| "Union[Dict[str, torch.Tensor], int]", |
| lhs["dict_keyword"]) |
| |
| self._assert_passes(template, |
| "Union[Dict[str, torch.Tensor], int]", |
| lhs["dict_keyword_with_iterable"]) |
| |
| self._assert_passes(template, |
| "Union[Dict[str, torch.Tensor], int]", |
| lhs["dict_keyword_with_empty_iterable"]) |
| |
| self._assert_passes(template, |
| "Union[Dict[str, torch.Tensor], int]", |
| lhs["dict_keyword_with_mapping"]) |
| |
| self._assert_passes(template, |
| "Union[Dict[str, torch.Tensor], int]", |
| lhs["dict_keyword_with_mapping_and_kwargs"]) |
| |
| # See above--string frontend does not support tuple unpacking |
| # self._assert_passes(template, |
| # "Union[Dict[str, torch.Tensor], int]", |
| # lhs["dict_keyword_with_internal_aggregate_function"]) |
| # |
| # self._assert_passes(template, |
| # "Union[Dict[str, torch.Tensor], int]", |
| # lhs["dict_comprehension_of_str_tensor"]) |
| |
| # self._assert_raises(template, |
| # "Union[Dict[str, torch.Tensor], int]", |
| # lhs["dict_comprehension_of_str_int"], |
| # "foobar") |
| |
| # self._assert_raises(template, |
| # "Union[Dict[str, torch.Tensor], int]", |
| # lhs["dict_comprehension_of_mixed"], |
| # "foobar") |
