code: def _unflatten_model_state_dict(
      model: nn.Module,
      state_dict: Union[
          Dict[nn.Module, Dict[str, ValueType]], Dict[str, ValueType]
      ],
  ) -> Dict[str, ValueType]:
      if not state_dict:
          return {}
  
      if isinstance(next(iter(state_dict.keys())), nn.Module):
          cast_state_dict = cast(
              Dict[nn.Module, Dict[str, ValueType]], state_dict
          )
          new_state_dict: Dict[str, ValueType] = {}
          for submodule, sub_state_dict in cast_state_dict.items():
              for name, m in model.named_modules():
                  if m != submodule:
                      continue
  
                  fqns = _get_fqns(model, name)
                  assert (
                      len(fqns) == 1
                  ), "FQNs for a submodule should only have 1 element"
                  prefix = f"{next(iter(fqns))}."
                  new_state_dict.update(
                      {
                          prefix + subfqn: value
                          for subfqn, value in sub_state_dict.items()
                      }
                  )
          return new_state_dict
      else:
          return cast(Dict[str, ValueType], state_dict)
  
imports: import contextlib
  import functools
  import gc
  from dataclasses import asdict, dataclass, field
  from itertools import chain
  from typing import Any, Callable, cast, Dict, Iterable, List, no_type_check, Optional, Set, Tuple, Union
  import torch
  import torch.distributed
  import torch.nn
  from torch.distributed._shard.sharded_tensor import ShardedTensor
  from torch.distributed._state_dict_utils import _gather_state_dict, _offload_state_dict_to_cpu
  from torch.distributed._tensor import DTensor
  from torch.distributed.fsdp import FullOptimStateDictConfig, FullStateDictConfig, FullyShardedDataParallel, OptimStateDictConfig, ShardedOptimStateDictConfig, ShardedStateDictConfig, StateDictConfig, StateDictType
  from torch.distributed.fsdp._common_utils import _get_module_fsdp_state_if_fully_sharded_module, FSDP_WRAPPED_MODULE
  from torch.nn.modules.module import _IncompatibleKeys
  from torch.nn.parallel import DistributedDataParallel
label_desc: Use image node_id 15 for calling a global function with example usage: _unflatten_model_state_dict(model, state_dict) and returns: dict, new_state_dict, cast