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))',
      }
  
      """
          Dict[str, torch.Tensor] | Dict[str, int]
          """
      self._assert_raises(
          template,
          "List[str] | List[torch.Tensor]",
          lhs["dict_literal_empty"],
          "Expected an Union type annotation with an "
          "inner Dict type",
      )
  
      self._assert_passes(
          template,
          "Dict[str, torch.Tensor] | Dict[str, int]",
          lhs["dict_literal_of_str_tensor"],
      )
  
      self._assert_passes(
          template,
          "Dict[str, torch.Tensor] | Dict[str, int]",
          lhs["dict_literal_of_str_int"],
      )
  
      self._assert_raises(
          template,
          "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, "Dict[str, torch.Tensor] | Dict[str, int]",
      #              lhs["dict_comprehension_of_str_tensor"])
  
      # self._assert_passes(template, "Dict[str, torch.Tensor] | Dict[str, int]",
      #              lhs["dict_comprehension_of_str_int"])
  
      # self._assert_raises(template, "Dict[str, torch.Tensor] | Dict[str, int]",
      #              lhs["dict_comprehension_of_mixed"],
      #              "foobar")
  
      # self._assert_passes(template,
      #                    "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,
          "Dict[str, torch.Tensor] | Dict[str, int]",
          lhs["dict_keyword"],
          "full type inference is not yet supported",
      )
  
      self._assert_raises(
          template,
          "Dict[str, torch.Tensor] | Dict[str, int]",
          lhs["dict_keyword_with_iterable"],
          "full type inference is not yet supported",
      )
  
      self._assert_raises(
          template,
          "Dict[str, torch.Tensor] | Dict[str, int]",
          lhs["dict_keyword_with_empty_iterable"],
          "full type inference is not yet supported",
      )
  
      self._assert_raises(
          template,
          "Dict[str, torch.Tensor] | Dict[str, int]",
          lhs["dict_keyword_with_mapping"],
          "full type inference is not yet supported",
      )
  
      self._assert_raises(
          template,
          "Dict[str, torch.Tensor] | Dict[str, int]",
          lhs["dict_keyword_with_mapping_and_kwargs"],
          "full type inference is not yet supported",
      )
  
      """
          int | torch.Tensor
          """
      self._assert_raises(
          template,
          "int | torch.Tensor",
          lhs["dict_literal_empty"],
          "Expected an Union type annotation with "
          "an inner Dict type",
      )
  
      self._assert_raises(
          template,
          "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, "int | torch.Tensor",
      #              lhs["dict_comprehension_of_tensor"],
      #              "foobar")
  
      """
          Dict[str, torch.Tensor] | int
          """
      self._assert_passes(
          template,
          "Dict[str, torch.Tensor] | int",
          lhs["dict_literal_empty"],
      )
  
      self._assert_passes(
          template,
          "Dict[str, torch.Tensor] | int",
          lhs["dict_literal_of_str_tensor"],
      )
  
      self._assert_raises(
          template,
          "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,
          "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, "Dict[str, torch.Tensor] | int", lhs["dict_keyword"]
      )
  
      self._assert_passes(
          template,
          "Dict[str, torch.Tensor] | int",
          lhs["dict_keyword_with_iterable"],
      )
  
      self._assert_passes(
          template,
          "Dict[str, torch.Tensor] | int",
          lhs["dict_keyword_with_empty_iterable"],
      )
  
      self._assert_passes(
          template,
          "Dict[str, torch.Tensor] | int",
          lhs["dict_keyword_with_mapping"],
      )
  
      self._assert_passes(
          template,
          "Dict[str, torch.Tensor] | int",
          lhs["dict_keyword_with_mapping_and_kwargs"],
      )
  
imports: import io
  import os
  import sys
  import unittest
  from enum import Enum
  from textwrap import dedent
  from typing import Dict, List, Optional, Tuple, Union
  import torch
  from torch.testing import FileCheck
  from torch.testing._internal.jit_utils import JitTestCase, make_global
label_desc: Use image node_id 52 for calling the TestUnion obj's underlying member method code with example usage: obj.test_union_with_dict_assignment() without return types