| 2240 | |
| 2241 | |
| 2242 | class ExportedProgramDeserializer: |
| 2243 | def __init__(self): |
| 2244 | pass |
| 2245 | |
| 2246 | def deserialize_range_constraints( |
| 2247 | self, |
| 2248 | symbol_name_to_range: Dict[str, symbolic_shapes.ValueRanges], |
| 2249 | symbol_name_to_symbol: Dict[str, sympy.Symbol], |
| 2250 | ) -> Dict[sympy.Symbol, ValueRanges]: |
| 2251 | range_constraints = {} |
| 2252 | for k, v in symbol_name_to_range.items(): |
| 2253 | if symbol := symbol_name_to_symbol.get(k): |
| 2254 | range_constraints[symbol] = v # type: ignore[arg-type] |
| 2255 | else: |
| 2256 | log.warning(f"Symbol {k} did not appear in the graph that was deserialized") # noqa: G004 |
| 2257 | return range_constraints |
| 2258 | |
| 2259 | def deserialize( |
| 2260 | self, |
| 2261 | exported_program: ExportedProgram, |
| 2262 | state_dict: Union[Dict[str, torch.Tensor], bytes], |
| 2263 | constants: Union[Dict[str, torch.Tensor], bytes], |
| 2264 | example_inputs: Optional[ |
| 2265 | Union[Tuple[Tuple[torch.Tensor, ...], Dict[str, Any]], bytes] |
| 2266 | ] = None, |
| 2267 | ) -> ep.ExportedProgram: |
| 2268 | assert isinstance(exported_program, ExportedProgram) |
| 2269 | version = exported_program.schema_version |
| 2270 | |
| 2271 | # TODO(zhxchen17) blocked on thrift schema refactor |
| 2272 | if version.major != SCHEMA_VERSION[0] and not ( |
| 2273 | version.major == 0 and version.minor == 0 |
| 2274 | ): |
| 2275 | raise SerializeError( |
| 2276 | f"Serialized schema version {exported_program.schema_version} " |
| 2277 | f"does not match our current schema version {SCHEMA_VERSION}." |
| 2278 | ) |
| 2279 | |
| 2280 | symbol_name_to_range = { |
| 2281 | k: symbolic_shapes.ValueRanges( |
| 2282 | _int_to_sympy_int(v.min_val), _int_to_sympy_int(v.max_val) |
| 2283 | ) |
| 2284 | for k, v in exported_program.range_constraints.items() |
| 2285 | } |
| 2286 | res = GraphModuleDeserializer().deserialize( |
| 2287 | exported_program.graph_module, |
| 2288 | state_dict, |
| 2289 | constants, |
| 2290 | example_inputs, |
| 2291 | symbol_name_to_range, |
| 2292 | ) |
| 2293 | range_constraints = self.deserialize_range_constraints( |
| 2294 | symbol_name_to_range, |
| 2295 | res.names_to_symbols, |
| 2296 | ) |
| 2297 | exported_program = ep.ExportedProgram( |
| 2298 | root=res.graph_module, |
| 2299 | graph=res.graph_module.graph, |