class_docstr: 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`.
class_docstr_tok: ["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","ca","n't","use","`","checkScript","`","in","our","test","cases","because","either","the","eager","mode","or","the","script","mode","would","n't","run","!","So",",","some","test","cases","have","separate","but","equivalent","functions","to","emulate","`","checkScript","`","."]
code: 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"],
      )
  
imports: 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
  from torch.testing._internal.jit_utils import JitTestCase, make_global
label_desc: Use image node_id 51 for calling the TestUnion obj's underlying member method code with example usage: obj.test_union_with_dict_assignment() without return types