code: def _shard_orig_param_state(
      fsdp_param_info: FSDPParamInfo,
      fqn: str,
      optim_state: Dict[str, Any],
  ) -> Dict[str, Any]:
      """
      Shard the optimizer state for the original parameter with the name ``fqn``.
      This API should only be used when ``use_orig_params`` is True.
      """
      if not optim_state:
          return {}
      fsdp_state = fsdp_param_info.state
      flat_param = fsdp_param_info.handle.flat_param
      param_idx = fsdp_param_info.param_indices[fqn]
      shard_param_info = flat_param._shard_param_infos[param_idx]  # type: ignore[attr-defined]
      optim_state = _gather_state_dict(
          optim_state,
          pg=fsdp_state.process_group,
          device=fsdp_state.compute_device,
      )
      if not shard_param_info.in_shard:
          return {}
      # Flatten and shard the state.
      new_optim_state: Dict[str, Any] = {}
      intra_param_start_idx = shard_param_info.intra_param_start_idx
      intra_param_end_idx = shard_param_info.intra_param_end_idx
      for state_name, value in optim_state.items():
          if (
              torch.is_tensor(value)
              and value.dim() > 0
              and fsdp_state.sharding_strategy
              != ShardingStrategy.NO_SHARD
          ):
              value = value.flatten()[intra_param_start_idx : intra_param_end_idx + 1].clone()  # type: ignore[operator]
          new_optim_state[state_name] = value
      return new_optim_state
  
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 8 for calling a global function with example usage: _shard_orig_param_state(fsdp_param_info, fqn, optim_state) and returns: new_optim_state, dict, dict