class_docstr: Takes a dict of parameters and exposes state_dict by parameter key.
  
  This implementation is much stricter than the one in torch.Optimizer:
  it requires implementations to fully initialize their state during first optimization iteration,
  and it prohibits loading an empty state into already initialized KeyedOptimizer and vise versa.
  
  It also doesn't expose param_groups in state_dict() by default
  Old behavior can be switch on by setting save_param_groups flag.
  The reason is that during distributed training not all parameters are present on all ranks
  and we identify param_group by its parameters.
  In addition to that, param_groups are typically re-set during training initialization,
  so it makes little sense to save them as a part of the state to begin with.
class_docstr_tok: ["Takes","a","dict","of","parameters","and","exposes","state_dict","by","parameter","key",".","This","implementation","is","much","stricter","than","the","one","in","torch.Optimizer",":","it","requires","implementations","to","fully","initialize","their","state","during","first","optimization","iteration",",","and","it","prohibits","loading","an","empty","state","into","already","initialized","KeyedOptimizer","and","vise","versa",".","It","also","does","n't","expose","param_groups","in","state_dict","(",")","by","default","Old","behavior","can","be","switch","on","by","setting","save_param_groups","flag",".","The","reason","is","that","during","distributed","training","not","all","parameters","are","present","on","all","ranks","and","we","identify","param_group","by","its","parameters",".","In","addition","to","that",",","param_groups","are","typically","re-set","during","training","initialization",",","so","it","makes","little","sense","to","save","them","as","a","part","of","the","state","to","begin","with","."]
docstr: Returned state and param_groups will contain parameter keys
  instead of parameter indices in torch.Optimizer.
  This allows for advanced functionality like optimizer re-sharding to be implemented.
  
  Can also handle classes and supported data structures that follow the PyTorch stateful
  protocol.
docstr_tok: ["Returned","state","and","param_groups","will","contain","parameter","keys","instead","of","parameter","indices","in","torch.Optimizer",".","This","allows","for","advanced","functionality","like","optimizer","re-sharding","to","be","implemented",".","Can","also","handle","classes","and","supported","data","structures","that","follow","the","PyTorch","stateful","protocol","."]
code: def state_dict(self) -> Dict[str, Any]:
      """
      Returned state and param_groups will contain parameter keys
      instead of parameter indices in torch.Optimizer.
      This allows for advanced functionality like optimizer re-sharding to be implemented.
  
      Can also handle classes and supported data structures that follow the PyTorch stateful
      protocol.
      """
  
      param_groups = self.param_groups
      params = self.params
      param_to_key = {param: key for key, param in params.items()}
  
      ret_state = {
          param_to_key[param]: self._extract_state_dict_content(
              param_state
          )
          for param, param_state in self.state.items()
      }
  
      ret_groups = []
      for group in param_groups:
          param_keys = []
          for param in group["params"]:
              param_keys.append(param_to_key[param])
          ret_group = {"params": sorted(param_keys)}
          for k, v in group.items():
              if k != "params":
                  ret_group[k] = deepcopy(v)
          ret_groups.append(ret_group)
  
      ret: Dict[str, object] = {"state": ret_state}
      if self.defaults["_save_param_groups"]:
          ret["param_groups"] = ret_groups
      return ret
  
imports: import json
  from copy import deepcopy
  from typing import Any, Callable, Collection, Dict, List, Mapping, Optional, OrderedDict, Set, Tuple, Union
  import torch
  from torch import optim
  from torchrec.distributed.types import ShardedTensor
label_desc: Use image node_id 4 for calling the KeyedOptimizer obj's underlying member method code with example usage: obj.state_dict() and returns: ret