(self, graph_module: fx.GraphModule)
| 648 | ) |
| 649 | |
| 650 | def call(self, graph_module: fx.GraphModule) -> PassResult: |
| 651 | if not getattr(self, "_initialized", False): |
| 652 | raise ExportPassBaseError( |
| 653 | "ExportPass is not initialized with __init__().", |
| 654 | ) |
| 655 | |
| 656 | inputs = self.inputs(graph_module) |
| 657 | |
| 658 | fake_tensor_mode = None |
| 659 | for i in inputs: |
| 660 | if isinstance(i, FakeTensor): |
| 661 | assert ( |
| 662 | fake_tensor_mode is None or fake_tensor_mode is i.fake_mode |
| 663 | ), "Multiple fake tensor mode detected." |
| 664 | fake_tensor_mode = i.fake_mode |
| 665 | if fake_tensor_mode is None: |
| 666 | fake_tensor_mode = FakeTensorMode(allow_non_fake_inputs=True) |
| 667 | dispatcher_mode = nullcontext() # type: ignore[assignment] |
| 668 | else: |
| 669 | fake_tensor_mode.allow_non_fake_inputs = True |
| 670 | dispatcher_mode = enable_python_dispatcher() # type: ignore[assignment] |
| 671 | self.tracer.fake_tensor_mode = fake_tensor_mode |
| 672 | self.fake_tensor_mode = fake_tensor_mode |
| 673 | |
| 674 | with fake_tensor_mode, dispatcher_mode: # type: ignore[assignment, union-attr] |
| 675 | result = self.call_submodule(graph_module, tuple(inputs)) |
| 676 | |
| 677 | return result |
| 678 | |
| 679 | |
| 680 | class ExportPass(_ExportPassBase): |
nothing calls this directly
no test coverage detected