text: # Owner(s): ["oncall: package/deploy"]
  
  from io import BytesIO
  from textwrap import dedent
  from unittest import skipIf
  
  import torch
  from torch.package import PackageExporter, PackageImporter
  from torch.testing._internal.common_utils import (
      IS_FBCODE,
      IS_SANDCASTLE,
      run_tests,
  )
  
  try:
      from .common import PackageTestCase
  except ImportError:
      # Support the case where we run this file directly.
      from common import PackageTestCase
  
  try:
      from torchvision.models import resnet18
  
      HAS_TORCHVISION = True
  except ImportError:
      HAS_TORCHVISION = False
  skipIfNoTorchVision = skipIf(not HAS_TORCHVISION, "no torchvision")
  
  
  class TestPackageScript(PackageTestCase):
      """Tests for compatibility with TorchScript."""
  
      def test_package_interface(self):
          """Packaging an interface class should work correctly."""
  
          import package_a.fake_interface as fake
  
          uses_interface = fake.UsesInterface()
          scripted = torch.jit.script(uses_interface)
          scripted.proxy_mod = torch.jit.script(fake.NewModule())
  
          buffer = BytesIO()
          with PackageExporter(buffer) as pe:
              pe.intern("**")
              pe.save_pickle("model", "model.pkl", uses_interface)
          buffer.seek(0)
  
          package_importer = PackageImporter(buffer)
          loaded = package_importer.load_pickle("model", "model.pkl")
  
          scripted_loaded = torch.jit.script(loaded)
          scripted_loaded.proxy_mod = torch.jit.script(fake.NewModule())
  
          input = torch.tensor(1)
  
          self.assertEqual(scripted(input), scripted_loaded(input))
  
      def test_different_package_interface(self):
          """Test a case where the interface defined in the package is
          different than the one defined in the loading environment, to make
          sure TorchScript can distinguish between the two.
          """
          # Import one version of the interface
          import package_a.fake_interface as fake
  
          # Simulate a package that contains a different version of the
          # interface, with the exact same name.
          buffer = BytesIO()
          with PackageExporter(buffer) as pe:
              pe.save_source_string(
                  fake.__name__,
                  dedent(
                      """\
                      import torch
                      from torch import Tensor
  
                      @torch.jit.interface
                      class ModuleInterface(torch.nn.Module):
                          def one(self, inp1: Tensor) -> Tensor:
                              pass
  
                      class ImplementsInterface(torch.nn.Module):
                          def one(self, inp1: Tensor) -> Tensor:
                              return inp1 + 1
  
                      class UsesInterface(torch.nn.Module):
                          proxy_mod: ModuleInterface
  
                          def __init__(self):
                              super().__init__()
                              self.proxy_mod = ImplementsInterface()
  
                          def forward(self, input: Tensor) -> Tensor:
                              return self.proxy_mod.one(input)
                      """
                  ),
              )
          buffer.seek(0)
  
          package_importer = PackageImporter(buffer)
          diff_fake = package_importer.import_module(fake.__name__)
          # We should be able to script successfully.
          torch.jit.script(diff_fake.UsesInterface())
  
      def test_package_script_class(self):
          import package_a.fake_script_class as fake
  
          buffer = BytesIO()
          with PackageExporter(buffer) as pe:
              pe.save_module(fake.__name__)
          buffer.seek(0)
  
          package_importer = PackageImporter(buffer)
          loaded = package_importer.import_module(fake.__name__)
  
          input = torch.tensor(1)
          self.assertTrue(
              torch.allclose(
                  fake.uses_script_class(input), loaded.uses_script_class(input)
              )
          )
  
      def test_different_package_script_class(self):
          """Test a case where the script class defined in the package is
          different than the one defined in the loading environment, to make
          sure TorchScript can distinguish between the two.
          """
          import package_a.fake_script_class as fake
  
          # Simulate a package that contains a different version of the
          # script class ,with the attribute `bar` instead of `foo`
          buffer = BytesIO()
          with PackageExporter(buffer) as pe2:
              pe2.save_source_string(
                  fake.__name__,
                  dedent(
                      """\
                      import torch
  
                      @torch.jit.script
                      class MyScriptClass:
                          def __init__(self, x):
                              self.bar = x
                      """
                  ),
              )
          buffer.seek(0)
  
          package_importer = PackageImporter(buffer)
          diff_fake = package_importer.import_module(fake.__name__)
          input = torch.rand(2, 3)
          loaded_script_class = diff_fake.MyScriptClass(input)
          orig_script_class = fake.MyScriptClass(input)
          self.assertEqual(loaded_script_class.bar, orig_script_class.foo)
  
      def test_save_scriptmodule(self):
          """
          Test basic saving of ScriptModule.
          """
          from package_a.test_module import ModWithTensor
  
          scripted_mod = torch.jit.script(ModWithTensor(torch.rand(1, 2, 3)))
  
          buffer = BytesIO()
          with PackageExporter(buffer) as e:
              e.save_pickle("res", "mod.pkl", scripted_mod)
  
          buffer.seek(0)
          importer = PackageImporter(buffer)
          loaded_mod = importer.load_pickle("res", "mod.pkl", map_location="cpu")
          input = torch.rand(1, 2, 3)
          self.assertEqual(loaded_mod(input), scripted_mod(input))
  
      @skipIf(
          IS_FBCODE or IS_SANDCASTLE,
          "Tests that use temporary files are disabled in fbcode",
      )
      def test_save_scriptmodule_file(self):
          """
          Test basic saving of ScriptModule in file.
          """
          from package_a.test_module import ModWithTensor
  
          scripted_mod = torch.jit.script(ModWithTensor(torch.rand(1, 2, 3)))
  
          filename = self.temp()
          with PackageExporter(filename) as e:
              e.save_pickle("res", "mod.pkl", scripted_mod)
  
          importer = PackageImporter(filename)
          loaded_mod = importer.load_pickle("res", "mod.pkl")
          input = torch.rand(1, 2, 3)
          self.assertEqual(loaded_mod(input), scripted_mod(input))
  
      def test_save_scriptmodule_with_submods(self):
          """
          Test basic saving of ScriptModule with submodule.
          """
          from package_a.test_module import ModWithSubmod, ModWithTensor
  
          scripted_mod = torch.jit.script(
              ModWithSubmod(ModWithTensor(torch.rand(1, 2, 3)))
          )
  
          buffer = BytesIO()
          with PackageExporter(buffer) as e:
              e.save_pickle("res", "mod.pkl", scripted_mod)
  
          buffer.seek(0)
          importer = PackageImporter(buffer)
          loaded_mod = importer.load_pickle("res", "mod.pkl", map_location="cpu")
          input = torch.rand(1, 2, 3)
          self.assertEqual(loaded_mod(input), scripted_mod(input))
  
      def test_save_scriptmodules_submod_redefinition(self):
          """
          Test to verify saving multiple ScriptModules with same top module
          but different submodules works. Submodule is redefined to between
          the defintion of the top module to check that the different concrete
          types of the modules are thoroughly recognized by serializaiton code.
          """
  
          class Submod(torch.nn.Module):
              def __init__(self):
                  super().__init__()
  
              def forward(self, input: str):
                  input = input + "_submod"
                  return input
  
          class TopMod(torch.nn.Module):
              def __init__(self):
                  super().__init__()
                  self.modB = Submod()
  
              def forward(self, input: str):
                  return self.modB(input)
  
          scripted_mod_0 = torch.jit.script(TopMod())
  
          # redefinition is intentional, change single inner string
          # string attribute, should trigger new module type
          class Submod(torch.nn.Module):  # noqa: F811
              def __init__(self):
                  super().__init__()
  
              def forward(self, input: str):
                  input = input + "_submod(changed)"
                  return input
  
          scripted_mod_1 = torch.jit.script(TopMod())
  
          buffer = BytesIO()
          with PackageExporter(buffer) as e:
              e.save_pickle("res", "mod1.pkl", scripted_mod_0)
              e.save_pickle("res", "mod2.pkl", scripted_mod_1)
  
          buffer.seek(0)
          importer = PackageImporter(buffer)
          loaded_mod_0 = importer.load_pickle("res", "mod1.pkl")
          loaded_mod_1 = importer.load_pickle("res", "mod2.pkl")
          self.assertEqual(loaded_mod_0("input"), scripted_mod_0("input"))
          self.assertEqual(loaded_mod_1("input"), scripted_mod_1("input"))
          self.assertNotEqual(loaded_mod_0("input"), loaded_mod_1("input"))
  
      def test_save_independent_scriptmodules(self):
          """
          Test to verify saving multiple ScriptModules with completely
          separate code works.
          """
          from package_a.test_module import ModWithTensor, SimpleTest
  
          scripted_mod_0 = torch.jit.script(SimpleTest())
          scripted_mod_1 = torch.jit.script(ModWithTensor(torch.rand(1, 2, 3)))
  
          buffer = BytesIO()
          with PackageExporter(buffer) as e:
              e.save_pickle("res", "mod1.pkl", scripted_mod_0)
              e.save_pickle("res", "mod2.pkl", scripted_mod_1)
  
          buffer.seek(0)
          importer = PackageImporter(buffer)
          loaded_mod_0 = importer.load_pickle("res", "mod1.pkl")
          loaded_mod_1 = importer.load_pickle("res", "mod2.pkl")
          input = torch.rand(1, 2, 3)
          self.assertEqual(loaded_mod_0(input), scripted_mod_0(input))
          self.assertEqual(loaded_mod_1(input), scripted_mod_1(input))
  
      def test_save_repeat_scriptmodules(self):
          """
          Test to verify saving multiple different modules and
          repeats of same scriptmodule in package works. Also tests that
          PyTorchStreamReader isn't having code hidden from
          PyTorchStreamWriter writing ScriptModule code files multiple times.
          """
          from package_a.test_module import (
              ModWithSubmodAndTensor,
              ModWithTensor,
              SimpleTest,
          )
  
          scripted_mod_0 = torch.jit.script(SimpleTest())
          scripted_mod_1 = torch.jit.script(ModWithTensor(torch.rand(1, 2, 3)))
          scripted_mod_2 = torch.jit.script(
              ModWithSubmodAndTensor(
                  torch.rand(1, 2, 3), ModWithTensor(torch.rand(1, 2, 3))
              )
          )
  
          buffer = BytesIO()
          with PackageExporter(buffer) as e:
              e.save_pickle("res", "mod0.pkl", scripted_mod_0)
              e.save_pickle("res", "mod1.pkl", scripted_mod_1)
              e.save_pickle("res", "mod2.pkl", scripted_mod_0)
              e.save_pickle("res", "mod3.pkl", scripted_mod_1)
              e.save_pickle("res", "mod4.pkl", scripted_mod_2)
  
          buffer.seek(0)
          importer = PackageImporter(buffer)
          loaded_mod_0 = importer.load_pickle("res", "mod0.pkl")
          loaded_mod_1 = importer.load_pickle("res", "mod3.pkl")
          loaded_mod_2 = importer.load_pickle("res", "mod4.pkl")
          input = torch.rand(1, 2, 3)
          self.assertEqual(loaded_mod_0(input), scripted_mod_0(input))
          self.assertEqual(loaded_mod_1(input), scripted_mod_1(input))
          self.assertEqual(loaded_mod_2(input), scripted_mod_2(input))
  
      def test_scriptmodules_repeat_save(self):
          """
          Test to verify saving and loading same ScriptModule object works
          across multiple packages.
          """
          from package_a.test_module import ModWithSubmodAndTensor, ModWithTensor
  
          scripted_mod_0 = torch.jit.script(ModWithTensor(torch.rand(1, 2, 3)))
          scripted_mod_1 = torch.jit.script(
              ModWithSubmodAndTensor(
                  torch.rand(1, 2, 3), ModWithTensor(torch.rand(1, 2, 3))
              )
          )
  
          buffer_0 = BytesIO()
          with PackageExporter(buffer_0) as e:
              e.save_pickle("res", "mod1.pkl", scripted_mod_0)
  
          buffer_0.seek(0)
          importer_0 = PackageImporter(buffer_0)
          loaded_module_0 = importer_0.load_pickle("res", "mod1.pkl")
  
          buffer_1 = BytesIO()
          with PackageExporter(buffer_1) as e:
              e.save_pickle("res", "mod1.pkl", scripted_mod_1)
              e.save_pickle("res", "mod2.pkl", loaded_module_0)
  
          buffer_1.seek(0)
          importer_1 = PackageImporter(buffer_1)
          loaded_module_1 = importer_1.load_pickle("res", "mod1.pkl")
          reloaded_module_0 = importer_1.load_pickle("res", "mod2.pkl")
  
          input = torch.rand(1, 2, 3)
          self.assertEqual(loaded_module_0(input), scripted_mod_0(input))
          self.assertEqual(loaded_module_0(input), reloaded_module_0(input))
          self.assertEqual(loaded_module_1(input), scripted_mod_1(input))
  
      @skipIfNoTorchVision
      def test_save_scriptmodule_only_necessary_code(self):
          """
          Test to verify when saving multiple packages with same CU
          that packages don't include unnecessary torchscript code files.
          The TorchVision code should only be saved in the package that
          relies on it.
          """
          from package_a.test_module import ModWithTensor
  
          class ModWithTorchVision(torch.nn.Module):
              def __init__(self, name: str):
                  super().__init__()
                  self.tvmod = resnet18()
  
              def forward(self, input):
                  return input * 4
  
          scripted_mod_0 = torch.jit.script(ModWithTorchVision("foo"))
          scripted_mod_1 = torch.jit.script(ModWithTensor(torch.rand(1, 2, 3)))
  
          buffer_0 = BytesIO()
          with PackageExporter(buffer_0) as e:
              e.save_pickle("res", "mod1.pkl", scripted_mod_0)
  
          buffer_0.seek(0)
          importer_0 = importer = PackageImporter(buffer_0)
  
          buffer_1 = BytesIO()
          with PackageExporter(buffer_1) as e:
              e.save_pickle("res", "mod1.pkl", scripted_mod_1)
  
          buffer_1.seek(0)
          importer_1 = PackageImporter(buffer_1)
  
          self.assertTrue("torchvision" in str(importer_0.file_structure()))
          self.assertFalse("torchvision" in str(importer_1.file_structure()))
  
      def test_save_scriptmodules_in_container(self):
          """
          Test saving of ScriptModules inside of container. Checks that relations
          between shared modules are upheld.
          """
          from package_a.test_module import ModWithSubmodAndTensor, ModWithTensor
  
          scripted_mod_a = torch.jit.script(ModWithTensor(torch.rand(1, 2, 3)))
          scripted_mod_b = torch.jit.script(
              ModWithSubmodAndTensor(torch.rand(1, 2, 3), scripted_mod_a)
          )
          script_mods_list = [scripted_mod_a, scripted_mod_b]
  
          buffer = BytesIO()
          with PackageExporter(buffer) as e:
              e.save_pickle("res", "list.pkl", script_mods_list)
  
          buffer.seek(0)
          importer = PackageImporter(buffer)
          loaded_mod_list = importer.load_pickle("res", "list.pkl")
          input = torch.rand(1, 2, 3)
          self.assertEqual(loaded_mod_list[0](input), scripted_mod_a(input))
          self.assertEqual(loaded_mod_list[1](input), scripted_mod_b(input))
  
      def test_save_eager_mods_sharing_scriptmodule(self):
          """
          Test saving of single ScriptModule shared by multiple
          eager modules (ScriptModule should be saved just once
          even though is contained in multiple pickles).
          """
          from package_a.test_module import ModWithSubmod, SimpleTest
  
          scripted_mod = torch.jit.script(SimpleTest())
  
          mod1 = ModWithSubmod(scripted_mod)
          mod2 = ModWithSubmod(scripted_mod)
  
          buffer = BytesIO()
          with PackageExporter(buffer) as e:
              e.intern("**")
              e.save_pickle("res", "mod1.pkl", mod1)
              e.save_pickle("res", "mod2.pkl", mod2)
  
          buffer.seek(0)
          importer = PackageImporter(buffer)
          file_structure = importer.file_structure()
          self.assertTrue(file_structure.has_file(".data/ts_code/0"))
          self.assertFalse(file_structure.has_file(".data/ts_code/1"))
  
      def test_load_shared_scriptmodules(self):
          """
          Test loading of single ScriptModule shared by multiple eager
          modules in single pickle (ScriptModule objects should be the same).
          """
          from package_a.test_module import (
              ModWithMultipleSubmods,
              ModWithSubmod,
              SimpleTest,
          )
  
          scripted_mod = torch.jit.script(SimpleTest())
  
          mod1 = ModWithSubmod(scripted_mod)
          mod2 = ModWithSubmod(scripted_mod)
  
          mod_parent = ModWithMultipleSubmods(mod1, mod2)
  
          buffer = BytesIO()
          with PackageExporter(buffer) as e:
              e.intern("**")
              e.save_pickle("res", "mod.pkl", mod_parent)
  
          buffer.seek(0)
          importer = PackageImporter(buffer)
  
          loaded_mod = importer.load_pickle("res", "mod.pkl")
          self.assertTrue(
              id(loaded_mod.mod1.script_mod) == id(loaded_mod.mod2.script_mod)
          )
  
      def test_save_shared_tensors(self):
          """
          Test tensors shared across eager and ScriptModules are serialized once.
          """
          from package_a.test_module import ModWithSubmodAndTensor, ModWithTensor
  
          shared_tensor = torch.rand(2, 3, 4)
          scripted_mod = torch.jit.script(ModWithTensor(shared_tensor))
  
          mod1 = ModWithSubmodAndTensor(shared_tensor, scripted_mod)
          mod2 = ModWithSubmodAndTensor(shared_tensor, scripted_mod)
  
          buffer = BytesIO()
          with PackageExporter(buffer) as e:
              e.intern("**")
              e.save_pickle("res", "tensor", shared_tensor)
              e.save_pickle("res", "mod1.pkl", mod1)
              e.save_pickle("res", "mod2.pkl", mod2)
  
          buffer.seek(0)
          importer = PackageImporter(buffer)
          loaded_mod_1 = importer.load_pickle("res", "mod1.pkl")
  
          # assert that there is only one storage stored in package
          file_structure = importer.file_structure(include=".data/*.storage")
          self.assertTrue(len(file_structure.children[".data"].children) == 1)
  
          input = torch.rand(2, 3, 4)
          self.assertEqual(loaded_mod_1(input), mod1(input))
  
      def test_load_shared_tensors(self):
          """
          Test tensors shared across eager and ScriptModules on load
          are the same.
          """
          from package_a.test_module import (
              ModWithTensor,
              ModWithTwoSubmodsAndTensor,
          )
  
          shared_tensor = torch.ones(3, 3)
  
          scripted_mod_0 = torch.jit.script(ModWithTensor(shared_tensor))
          scripted_mod_1 = torch.jit.script(ModWithTensor(shared_tensor))
  
          mod1 = ModWithTwoSubmodsAndTensor(shared_tensor, scripted_mod_0, scripted_mod_1)
  
          self.assertEqual(
              shared_tensor.storage()._cdata,
              scripted_mod_0.tensor.storage()._cdata,
          )
          self.assertEqual(
              shared_tensor.storage()._cdata,
              scripted_mod_1.tensor.storage()._cdata,
          )
  
          buffer = BytesIO()
          with PackageExporter(buffer) as e:
              e.intern("**")
              e.save_pickle("res", "mod1.pkl", mod1)
  
          buffer.seek(0)
          importer = PackageImporter(buffer)
          loaded_mod_1 = importer.load_pickle("res", "mod1.pkl")
  
          self.assertEqual(
              loaded_mod_1.tensor.storage()._cdata,
              loaded_mod_1.sub_mod_0.tensor.storage()._cdata,
          )
          self.assertEqual(
              loaded_mod_1.tensor.storage()._cdata,
              loaded_mod_1.sub_mod_1.tensor.storage()._cdata,
          )
  
          loaded_mod_1.tensor.add_(torch.ones(3, 3))
  
          self.assertTrue(
              torch.allclose(loaded_mod_1.tensor, loaded_mod_1.sub_mod_0.tensor)
          )
          self.assertTrue(
              torch.allclose(loaded_mod_1.tensor, loaded_mod_1.sub_mod_1.tensor)
          )
  
      def test_load_shared_tensors_repackaged(self):
          """
          Test tensors shared across eager and ScriptModules on load
          are the same across multiple package saves and loads. This is
          an important test because not all of the tensor information is restored
          in python between packages. The python identity is not maintained, but
          the backing cpp TensorImpl is. We load/save storages based off of this
          cpp TensorImpl and not the python identity.
          """
          from package_a.test_module import (
              ModWithTensor,
              ModWithTwoSubmodsAndTensor,
          )
  
          shared_tensor = torch.ones(3, 3)
  
          scripted_mod_0 = torch.jit.script(ModWithTensor(shared_tensor))
          scripted_mod_1 = torch.jit.script(ModWithTensor(shared_tensor))
  
          mod1 = ModWithTwoSubmodsAndTensor(shared_tensor, scripted_mod_0, scripted_mod_1)
  
          buffer_0 = BytesIO()
          with PackageExporter(buffer_0) as e:
              e.intern("**")
              e.save_pickle("res", "mod1.pkl", mod1)
  
          buffer_0.seek(0)
          importer_0 = PackageImporter(buffer_0)
          loaded_mod_0 = importer_0.load_pickle("res", "mod1.pkl")
  
          buffer_1 = BytesIO()
          with PackageExporter(buffer_1, importer=importer_0) as e:
              e.intern("**")
              e.save_pickle("res", "mod1.pkl", loaded_mod_0)
  
          buffer_1.seek(0)
          importer = PackageImporter(buffer_1)
          loaded_mod_1 = importer.load_pickle("res", "mod1.pkl")
  
          self.assertEqual(
              loaded_mod_1.tensor.storage()._cdata,
              loaded_mod_1.sub_mod_0.tensor.storage()._cdata,
          )
          self.assertEqual(
              loaded_mod_1.tensor.storage()._cdata,
              loaded_mod_1.sub_mod_1.tensor.storage()._cdata,
          )
  
          loaded_mod_1.tensor.add_(
              torch.ones(3, 3)
          )  # all tensors should reflect this change
  
          self.assertTrue(
              torch.allclose(loaded_mod_1.tensor, loaded_mod_1.sub_mod_0.tensor)
          )
          self.assertTrue(
              torch.allclose(loaded_mod_1.tensor, loaded_mod_1.sub_mod_1.tensor)
          )
  
      def test_saving_and_scripting_packaged_mod(self):
          """
          Test scripting a module loaded from a package
          and saving it in a new package as a script object.
          """
          from package_a.test_module import SimpleTest
  
          orig_mod = SimpleTest()
  
          buffer_0 = BytesIO()
          with PackageExporter(buffer_0) as e:
              e.intern("**")
              e.save_pickle("model", "model.pkl", orig_mod)
  
          buffer_0.seek(0)
          importer_0 = PackageImporter(buffer_0)
          loaded_mod = importer_0.load_pickle("model", "model.pkl")
  
          input = torch.rand(2, 3)
          self.assertEqual(loaded_mod(input), orig_mod(input))
  
          scripted_mod = torch.jit.script(loaded_mod)
  
          buffer_1 = BytesIO()
          with PackageExporter(buffer_1, importer=importer_0) as e:
              e.intern("**")
              e.save_pickle("res", "scripted_mod.pkl", scripted_mod)
  
          buffer_1.seek(0)
          importer_1 = PackageImporter(buffer_1)
          loaded_mod_scripted = importer_1.load_pickle("res", "scripted_mod.pkl")
  
          self.assertEqual(loaded_mod_scripted(input), orig_mod(input))
  
      def test_mixing_packaged_and_inline_modules(self):
          """
          Test saving inline and imported modules in same package with
          independent code.
          """
  
          class InlineMod(torch.nn.Module):
              def __init__(self, name: str):
                  super().__init__()
                  self.name = name
                  self.tensor = torch.rand(1, 2, 3)
  
              def forward(self, input: str):
                  input = input + "_modInline:" + self.name
                  return input, (self.tensor * 4)
  
          inline_mod = InlineMod("inline")
          scripted_inline = torch.jit.script(inline_mod)
  
          from package_a.test_module import SimpleTest
  
          imported_mod = SimpleTest()
          scripted_imported = torch.jit.script(imported_mod)
  
          buffer = BytesIO()
          with PackageExporter(buffer) as e:
              e.save_pickle("model", "inline.pkl", scripted_inline)
              e.save_pickle("model", "imported.pkl", scripted_imported)
  
          buffer.seek(0)
          importer = PackageImporter(buffer)
          loaded_inline = importer.load_pickle("model", "inline.pkl")
          loaded_imported = importer.load_pickle("model", "imported.pkl")
  
          input = torch.rand(2, 3)
          self.assertEqual(loaded_imported(input), imported_mod(input))
          self.assertEqual(loaded_inline("input"), inline_mod("input"))
  
      @skipIfNoTorchVision
      def test_mixing_packaged_and_inline_modules_shared_code(self):
          """
          Test saving inline and imported modules in same package that
          share code.
          """
  
          class TorchVisionTestInline(torch.nn.Module):
              def __init__(self):
                  super().__init__()
                  self.tvmod = resnet18()
  
              def forward(self, x):
                  x = a_non_torch_leaf(x, x)
                  return torch.relu(x + 3.0)
  
          def a_non_torch_leaf(a, b):
              return a + b
  
          inline_mod = TorchVisionTestInline()
          scripted_inline = torch.jit.script(inline_mod)
  
          from package_c.test_module import TorchVisionTest
  
          imported_mod = TorchVisionTest()
          scripted_imported = torch.jit.script(imported_mod)
  
          buffer = BytesIO()
          with PackageExporter(buffer) as e:
              e.save_pickle("model", "inline.pkl", scripted_inline)
              e.save_pickle("model", "imported.pkl", scripted_imported)
  
          buffer.seek(0)
          importer = PackageImporter(buffer)
          loaded_inline = importer.load_pickle("model", "inline.pkl")
          loaded_imported = importer.load_pickle("model", "imported.pkl")
  
          input = torch.rand(2, 3)
          self.assertEqual(loaded_imported(input), imported_mod(input))
          self.assertEqual(loaded_inline(input), inline_mod(input))
  
      def test_tensor_sharing_pickle(self):
          """Test that saving a ScriptModule and a separately saving a tensor
          object causes no issues.
          """
  
          class M(torch.nn.Module):
              def __init__(self):
                  super().__init__()
                  self.foo = torch.ones(2, 3)
  
              def forward(self):
                  return self.foo
  
          scripted_m = torch.jit.script(M())
          original_tensor = torch.ones(0)
  
          f = BytesIO()
          with torch.package.PackageExporter(f) as exporter:
              exporter.save_pickle("model", "model.pkl", scripted_m)
              exporter.save_pickle("model", "input.pkl", original_tensor)
  
          f.seek(0)
          # Should be able to load correctly
          importer = PackageImporter(f)
          loaded_m = importer.load_pickle("model", "model.pkl")
          loaded_tensor = importer.load_pickle("model", "input.pkl")
  
          self.assertEqual(scripted_m.foo, loaded_m.foo)
          self.assertEqual(original_tensor, loaded_tensor)
  
  
  if __name__ == "__main__":
      run_tests()
  