(
self,
state_dict: Dict[str, Dict[str, Union[Dict[str, Tensor], str]]],
*,
warmup_with_run: bool = False,
)
| 197 | return graph |
| 198 | |
| 199 | def load_runtime_state_dict( |
| 200 | self, |
| 201 | state_dict: Dict[str, Dict[str, Union[Dict[str, Tensor], str]]], |
| 202 | *, |
| 203 | warmup_with_run: bool = False, |
| 204 | ) -> None: |
| 205 | graph_dict = dict() |
| 206 | for _, sub_state_dict in state_dict.items(): |
| 207 | cache_order = sub_state_dict["cache_order"] |
| 208 | graph_dict[cache_order] = sub_state_dict |
| 209 | |
| 210 | if self._cache is None: |
| 211 | self._cache = LRUCache(self._cache_size) |
| 212 | for _, sub_state_dict in sorted(graph_dict.items()): |
| 213 | cache_key = sub_state_dict["cache_key"] |
| 214 | graph = self._cache.get(cache_key) |
| 215 | assert graph is None |
| 216 | graph = self._init_and_get_a_graph_in_cache(cache_key) |
| 217 | with AvoidRecursiveCacheCall(graph): |
| 218 | graph.load_runtime_state_dict( |
| 219 | sub_state_dict, warmup_with_run=warmup_with_run |
| 220 | ) |
| 221 | |
| 222 | def gen_key(self, *args, **kwargs): |
| 223 | flattened_shapes = [] |
nothing calls this directly
no test coverage detected