code: def _unflatten_orig_param_states(
      fsdp_param_info: FSDPParamInfo,
      output_states: Dict[str, Dict[str, Any]],
      state_name: str,
      shard_state: bool,
      to_save: bool,
      cpu_offload: bool,
  ) -> None:
      """
      Given a output state dict, ``output_states``, which the keys are FQNs to the
      original parameters (not FlatParameters nor parmeter ID), and the values
      are gathered states, unflatten the states to the original dimensions.
  
      This function performs the unflattening process in-place.
      """
      if not to_save:
          return
      flat_param = fsdp_param_info.handle.flat_param
      fsdp_state = fsdp_param_info.state
      for fqn, gathered_state in output_states.items():
          value = gathered_state[state_name]
          param_idx = fsdp_param_info.param_indices[fqn]
  
          # TODO: This solution is not general and only apply to PTD TP solution.
          if isinstance(value, DTensor):
              placement = value.placements[0]
              # If gathered state is a DTensor and its TP placement is not Replicate(), we need to
              # gather the tensor on its TP dimension before chunking them into DTensor again.
              if placement != Replicate():
                  placement_dim = placement.dim  # type: ignore[attr-defined]
                  value_local = value.redistribute(
                      placements=(Replicate(),)
                  )
                  reshape_size = list(flat_param._shapes[param_idx])
                  reshape_size[placement_dim] *= value.device_mesh.size(
                      0
                  )
                  reshape_size = torch.Size(reshape_size)
                  value = value.reshape(reshape_size)
              # If gathered state is a replicate DTensor, we directly reshape it.
              else:
                  value = value.reshape(flat_param._shapes[param_idx])
          else:
              # If gathered state is a tensor, we directly reshape it into unflatten state.
              value = value.reshape(flat_param._shapes[param_idx])
  
          if shard_state:
              osd_config = fsdp_state._optim_state_dict_config
              if getattr(osd_config, "_use_dtensor", False):
                  assert fsdp_state._device_mesh is not None
                  value = _ext_chunk_dtensor(
                      value,
                      fsdp_state.rank,
                      fsdp_state._device_mesh,
                      fsdp_state._fsdp_extension,
                  )
              else:
                  assert fsdp_state.process_group is not None
                  value = _ext_chunk_tensor(
                      value,
                      fsdp_state.rank,
                      fsdp_state.world_size,
                      fsdp_state._device_handle.device_count(),
                      fsdp_state.process_group,
                      fsdp_state._fsdp_extension,
                  )
          elif not cpu_offload:
              with SimpleProfiler.profile("clone"):
                  value = value.detach().clone()
  
          if cpu_offload:
              with SimpleProfiler.profile(SimpleProfiler.Type.D2H):
                  value = value.cpu()
          gathered_state[state_name] = value
  
imports: import copy
  import functools
  import logging
  import warnings
  from contextlib import ExitStack
  from dataclasses import dataclass, field
  from typing import Any, cast, Dict, Iterable, Iterator, List, NamedTuple, no_type_check, Optional, Sequence, Set, Tuple, Union
  import torch
  import torch.distributed
  import torch.distributed.fsdp._traversal_utils
  import torch.nn
  from torch.distributed._shard.sharded_tensor import ShardedTensor
  from torch.distributed._state_dict_utils import _gather_state_dict
  from torch.distributed._tensor import DTensor, Replicate
  from torch.distributed.distributed_c10d import _get_pg_default_device
  from torch.distributed.fsdp._common_utils import _apply_to_modules, _FSDPState, _get_module_fsdp_state_if_fully_sharded_module, _get_param_to_fqns, _module_handle, _named_parameters_with_duplicates, clean_tensor_name
  from torch.distributed.fsdp._debug_utils import SimpleProfiler
  from torch.distributed.fsdp._flat_param import FlatParameter, FlatParamHandle
  from torch.distributed.fsdp._fsdp_extensions import _ext_chunk_dtensor, _ext_chunk_tensor
  from torch.distributed.fsdp._runtime_utils import _lazy_init, _reset_flat_param_grad_info_if_needed
  from torch.distributed.fsdp.api import ShardingStrategy, StateDictSettings, StateDictType
  from torch.utils._pytree import tree_map_only
label_desc: Use image node_id 26 for calling a global function with example usage: _unflatten_orig_param_states(fsdp_param_info, output_states, state_name, shard_state, to_save, cpu_offload) without return types