class_docstr: A suite of tests for profile directed typing in TorchScript.
class_docstr_tok: ["A","suite","of","tests","for","profile","directed","typing","in","TorchScript","."]
code: 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),
      )
  
imports: import os
  import sys
  import torch
  from torch.testing._internal.jit_utils import JitTestCase, make_global
  from torch.jit._monkeytype_config import _IS_MONKEYTYPE_INSTALLED
  from typing import List, Dict, Tuple, Any, Optional, NamedTuple
  from torch.testing._internal.common_utils import NoTest
label_desc: Use image node_id 5 for calling the TestPDT obj's underlying member method code with example usage: obj.test_nn_module_with_export_function() and returns: unknown, unknown, unknown