code: def original_state_dict_hooks(model: Any):
      """
      Use this patch if you want to save/load state dict in the original state dict hierarchy.
  
      For example, when you already have a state dict for the base model / search space (which often
      happens when you have trained a supernet with one-shot strategies), the state dict isn't organized
      in the same way as when a sub-model is sampled from the search space. This patch will help
      the modules in the sub-model find the corresponding module in the base model.
  
      The code looks like,
  
      .. code-block:: python
  
          with original_state_dict_hooks(model):
              model.load_state_dict(state_dict_from_supernet, strict=False)  # supernet has extra keys
  
      Or vice-versa,
  
      .. code-block:: python
  
          with original_state_dict_hooks(model):
              supernet_style_state_dict = model.state_dict()
      """
  
      import torch.utils.hooks
      import torch.nn as nn
  
      assert isinstance(
          model, nn.Module
      ), "PyTorch is the only supported framework for now."
  
      # the following are written for pytorch only
  
      # first get the full mapping
      full_mapping = {}
  
      def full_mapping_in_module(src_prefix, tar_prefix, module):
          if hasattr(module, STATE_DICT_PY_MAPPING):
              # only values are complete
              local_map = getattr(module, STATE_DICT_PY_MAPPING)
          elif hasattr(module, STATE_DICT_PY_MAPPING_PARTIAL):
              # keys and values are both incomplete
              local_map = getattr(module, STATE_DICT_PY_MAPPING_PARTIAL)
              local_map = {
                  k: tar_prefix + v for k, v in local_map.items()
              }
          else:
              # no mapping
              local_map = {}
  
          if "__self__" in local_map:
              # special case, overwrite prefix
              tar_prefix = local_map["__self__"] + "."
  
          for key, value in local_map.items():
              if (
                  key != "" and key not in module._modules
              ):  # not a sub-module, probably a parameter
                  full_mapping[src_prefix + key] = value
  
          if src_prefix != tar_prefix:  # To deal with leaf nodes.
              for name, value in itertools.chain(
                  module._parameters.items(), module._buffers.items()
              ):  # direct children
                  if (
                      value is None
                      or name in module._non_persistent_buffers_set
                  ):
                      # it won't appear in state dict
                      continue
                  if (src_prefix + name) not in full_mapping:
                      full_mapping[src_prefix + name] = (
                          tar_prefix + name
                      )
  
          for name, child in module.named_children():
              # sub-modules
              full_mapping_in_module(
                  src_prefix + name + ".",
                  local_map.get(name, tar_prefix + name)
                  + ".",  # if mapping doesn't exist, respect the prefix
                  child,
              )
  
      full_mapping_in_module("", "", model)
  
      def load_state_dict_hook(
          state_dict,
          prefix,
          local_metadata,
          strict,
          missing_keys,
          unexpected_keys,
          error_msgs,
      ):
          reverse_mapping = defaultdict(list)
          for src, tar in full_mapping.items():
              reverse_mapping[tar].append(src)
  
          transf_state_dict = {}
          for src, tar_keys in reverse_mapping.items():
              if src in state_dict:
                  value = state_dict.pop(src)
                  for tar in tar_keys:
                      transf_state_dict[tar] = value
              else:
                  missing_keys.append(src)
          state_dict.update(transf_state_dict)
  
      def state_dict_hook(module, destination, prefix, local_metadata):
          result = {}
          for src, tar in full_mapping.items():
              if src in destination:
                  result[tar] = destination.pop(src)
              else:
                  raise KeyError(
                      f'"{src}" not in state dict, but found in mapping.'
                  )
          destination.update(result)
  
      hooks: List[torch.utils.hooks.RemovableHandle] = []
      try:
          hooks.append(
              model._register_load_state_dict_pre_hook(
                  load_state_dict_hook
              )
          )
          hooks.append(model._register_state_dict_hook(state_dict_hook))
          yield
      finally:
          for hook in hooks:
              hook.remove()
  
file_path: nni/nni/nas/utils/misc.py
imports: import inspect
  import itertools
  import warnings
  from collections import defaultdict
  from contextlib import contextmanager
  from typing import Any, List, cast
  from pathlib import Path
label_desc: Use image node_id 4 for calling a global function with example usage: original_state_dict_hooks(model) without return types