A subclass of nn.Module that is generated for modules containing delegated functions. This is can be created by calling `to_backend`.
| 52 | |
| 53 | |
| 54 | class LoweredBackendModule(torch.nn.Module): |
| 55 | """ |
| 56 | A subclass of nn.Module that is generated for modules containing |
| 57 | delegated functions. This is can be created by calling `to_backend`. |
| 58 | """ |
| 59 | |
| 60 | _backend_id: str # The backend's name |
| 61 | _processed_bytes: bytes # The delegate blobs created from backend.preprocess |
| 62 | _compile_specs: List[ |
| 63 | CompileSpec |
| 64 | ] # A list of backend-specific objects with static metadata to configure the "compilation" process. |
| 65 | _original_exported_program: ExportedProgram # The original EXIR module |
| 66 | _named_data_store_output: Optional[ |
| 67 | NamedDataStoreOutput |
| 68 | ] # Named Data serialized by the backend |
| 69 | meta: Optional[Dict[str, Any]] # Metadata for the lowered module |
| 70 | |
| 71 | def __init__( |
| 72 | self, |
| 73 | edge_program: ExportedProgram, |
| 74 | backend_id: str, |
| 75 | processed_bytes: bytes, |
| 76 | compile_specs: List[CompileSpec], |
| 77 | named_data_store_output: Optional[NamedDataStoreOutput] = None, |
| 78 | ) -> None: |
| 79 | super().__init__() |
| 80 | self._original_exported_program = edge_program |
| 81 | self._backend_id = backend_id |
| 82 | self._processed_bytes = processed_bytes |
| 83 | self._compile_specs = compile_specs |
| 84 | self._named_data_store_output = named_data_store_output |
| 85 | self.meta = None |
| 86 | |
| 87 | # pyre-ignore |
| 88 | def __deepcopy__(self, memo: Optional[Dict[int, Any]]) -> "LoweredBackendModule": |
| 89 | # Copy exported program |
| 90 | copied_program = ExportedProgram( |
| 91 | root=copy.deepcopy(self._original_exported_program.graph_module), |
| 92 | graph=copy.deepcopy(self._original_exported_program.graph), |
| 93 | graph_signature=copy.deepcopy( |
| 94 | self._original_exported_program.graph_signature |
| 95 | ), |
| 96 | state_dict=self._original_exported_program.state_dict, |
| 97 | range_constraints=copy.deepcopy( |
| 98 | self._original_exported_program.range_constraints |
| 99 | ), |
| 100 | module_call_graph=copy.deepcopy( |
| 101 | self._original_exported_program.module_call_graph |
| 102 | ), |
| 103 | constants=self._original_exported_program.constants, |
| 104 | verifiers=[copy.deepcopy(self._original_exported_program.verifier)], |
| 105 | ) |
| 106 | |
| 107 | res = LoweredBackendModule( |
| 108 | edge_program=copied_program, |
| 109 | backend_id=self._backend_id, |
| 110 | processed_bytes=self._processed_bytes, |
| 111 | compile_specs=copy.deepcopy(self._compile_specs, memo), |
no outgoing calls