docstr: Test that a type annotation can be provided for a ModuleDict that allows
  non-static indexing.
docstr_tok: ["Test","that","a","type","annotation","can","be","provided","for","a","ModuleDict","that","allows","non-static","indexing","."]
code: def test_typed_module_dict(self):
      """
      Test that a type annotation can be provided for a ModuleDict that allows
      non-static indexing.
      """
  
      @torch.jit.interface
      class ModuleInterface(torch.nn.Module):
          def forward(self, inp: Any) -> Any:
              pass
  
      class ImplementsInterface(torch.nn.Module):
          def forward(self, inp: Any) -> Any:
              if isinstance(inp, torch.Tensor):
                  return torch.max(inp, dim=0)
  
              return inp
  
      class DoesNotImplementInterface(torch.nn.Module):
          def forward(
              self, inp: torch.Tensor
          ) -> Tuple[torch.Tensor, torch.Tensor]:
              return torch.max(inp, dim=0)
  
      # Test annotation of submodule.
      class Mod(torch.nn.Module):
          def __init__(self):
              super().__init__()
              self.d = torch.nn.ModuleDict(
                  {"module": ImplementsInterface()}
              )
  
          def forward(self, x: torch.Tensor, key: str) -> Any:
              value: ModuleInterface = self.d[key]
              return value.forward(x)
  
      m = Mod()
      self.checkModule(m, (torch.randn(2, 2), "module"))
  
      # Test annotation of self.
      class ModDict(torch.nn.ModuleDict):
          def __init__(self):
              super().__init__({"module": ImplementsInterface()})
  
          def forward(self, x: torch.Tensor, key: str) -> Any:
              submodule: ModuleInterface = self[key]
              return submodule.forward(x)
  
      m = ModDict()
      self.checkModule(m, (torch.randn(2, 2), "module"))
  
      # Test error message thrown when annotated attribute does not comply with the
      # annotation.
      class ModWithWrongAnnotation(torch.nn.ModuleDict):
          def __init__(self):
              super().__init__()
              self.d = torch.nn.ModuleDict(
                  {"module": DoesNotImplementInterface()}
              )
  
          def forward(self, x: torch.Tensor, key: str) -> Any:
              submodule: ModuleInterface = self.d[key]
              return submodule.forward(x)
  
      with self.assertRaisesRegexWithHighlight(
          RuntimeError,
          r"Attribute module is not of annotated type",
          "self.d[key]",
      ):
          torch.jit.script(ModWithWrongAnnotation())
  
imports: import os
  import sys
  from typing import Any, List, Tuple
  from collections import OrderedDict
  import torch
  import torch.nn
  from torch.testing._internal.jit_utils import JitTestCase
label_desc: Use image node_id 12 for calling the TestModuleContainers obj's underlying member method code with example usage: obj.test_typed_module_dict() and returns: inp, torch, value, submodule, submodule, torch