(
self,
serialized_graph_module: GraphModule,
serialized_state_dict: Union[Dict[str, torch.Tensor], bytes],
constants: Union[Dict[str, Any], bytes],
example_inputs: Optional[
Union[Tuple[Tuple[torch.Tensor, ...], Dict[str, Any]], bytes]
] = None,
symbol_name_to_range: Optional[Dict[str, symbolic_shapes.ValueRanges]] = None,
)
| 1831 | ) |
| 1832 | |
| 1833 | def deserialize( |
| 1834 | self, |
| 1835 | serialized_graph_module: GraphModule, |
| 1836 | serialized_state_dict: Union[Dict[str, torch.Tensor], bytes], |
| 1837 | constants: Union[Dict[str, Any], bytes], |
| 1838 | example_inputs: Optional[ |
| 1839 | Union[Tuple[Tuple[torch.Tensor, ...], Dict[str, Any]], bytes] |
| 1840 | ] = None, |
| 1841 | symbol_name_to_range: Optional[Dict[str, symbolic_shapes.ValueRanges]] = None, |
| 1842 | ) -> Result: |
| 1843 | global _CURRENT_DESERIALIZER |
| 1844 | current_deserializer_state = _CURRENT_DESERIALIZER.copy() |
| 1845 | _CURRENT_DESERIALIZER.append(self) |
| 1846 | try: |
| 1847 | self.shape_env = symbolic_shapes.ShapeEnv(assume_static_by_default=True) |
| 1848 | self.fake_tensor_mode = FakeTensorMode( |
| 1849 | allow_fallback_kernels=False, |
| 1850 | allow_non_fake_inputs=True, |
| 1851 | shape_env=self.shape_env, |
| 1852 | ) |
| 1853 | self.symbol_name_to_symbol: Dict[str, sympy.Symbol] = {} |
| 1854 | self.constants = deserialize_torch_artifact(constants) |
| 1855 | self.signature = self.deserialize_signature( |
| 1856 | serialized_graph_module.signature |
| 1857 | ) |
| 1858 | |
| 1859 | # deserialization does analysis with checks on 0/1, so we create fake range constraints and |
| 1860 | # restore the original range constraints afterwards |
| 1861 | self.symbol_name_to_range = {} |
| 1862 | if symbol_name_to_range: |
| 1863 | for k, vr in symbol_name_to_range.items(): |
| 1864 | if math.isinf(vr.lower) and vr.lower < 0: |
| 1865 | lower = -math.inf |
| 1866 | elif math.isinf(vr.lower): |
| 1867 | lower = math.inf |
| 1868 | else: |
| 1869 | lower = int(vr.lower) |
| 1870 | |
| 1871 | if vr.upper >= 2: # max is >= 2, not sym bool range |
| 1872 | lower = max(2, lower) |
| 1873 | self.symbol_name_to_range[k] = symbolic_shapes.ValueRanges( |
| 1874 | _int_to_sympy_int(lower), vr.upper |
| 1875 | ) |
| 1876 | |
| 1877 | if example_inputs is not None and len(example_inputs) > 0: |
| 1878 | self.example_inputs = deserialize_torch_artifact(example_inputs) |
| 1879 | else: |
| 1880 | self.example_inputs = None |
| 1881 | self.deserialize_graph(serialized_graph_module.graph) |
| 1882 | |
| 1883 | module_call_graph = self.deserialize_module_call_graph( |
| 1884 | serialized_graph_module.module_call_graph |
| 1885 | ) |
| 1886 | return GraphModuleDeserializer.Result( |
| 1887 | graph_module=ep._create_graph_module_for_export( |
| 1888 | self.module, self.graph |
| 1889 | ), |
| 1890 | signature=self.signature, |
no test coverage detected