class_docstr: Tests for compatibility with FX.
class_docstr_tok: ["Tests","for","compatibility","with","FX","."]
code: def test_package_fx_custom_tracer(self):
      from package_a.test_all_leaf_modules_tracer import (
          TestAllLeafModulesTracer,
      )
      from package_a.test_module import (
          ModWithTwoSubmodsAndTensor,
          SimpleTest,
      )
  
      class SpecialGraphModule(torch.fx.GraphModule):
          def __init__(self, root, graph, info):
              super().__init__(root, graph)
              self.info = info
  
      sub_module = SimpleTest()
      module = ModWithTwoSubmodsAndTensor(
          torch.ones(3),
          sub_module,
          sub_module,
      )
      tracer = TestAllLeafModulesTracer()
      graph = tracer.trace(module)
  
      self.assertEqual(graph._tracer_cls, TestAllLeafModulesTracer)
  
      gm = SpecialGraphModule(module, graph, "secret")
      self.assertEqual(gm._tracer_cls, TestAllLeafModulesTracer)
  
      f = BytesIO()
      with PackageExporter(f) as pe:
          pe.intern("**")
          pe.save_pickle("model", "model.pkl", gm)
      f.seek(0)
  
      pi = PackageImporter(f)
      loaded_gm = pi.load_pickle("model", "model.pkl")
      self.assertEqual(
          type(loaded_gm).__class__.__name__,
          SpecialGraphModule.__class__.__name__,
      )
      self.assertEqual(loaded_gm.info, "secret")
  
      input_x = torch.randn(3)
      self.assertEqual(loaded_gm(input_x), gm(input_x))
  
file_path: pytorch/test/package/test_package_fx.py
imports: from io import BytesIO
  import torch
  from torch.fx import Graph, GraphModule, symbolic_trace
  from torch.package import ObjMismatchError, PackageExporter, PackageImporter, sys_importer
  from torch.testing._internal.common_utils import run_tests
label_desc: Use image node_id 5 for calling the TestPackageFX obj's underlying member method code with example usage: obj.test_package_fx_custom_tracer() without return types