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: 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.
  
  Because of introduced strictness it allows us to:
      * do compatibility checks for state and param_groups, which improves usability
      * avoid state duplication by directly copying into state tensors, e.g.
        optimizer.step()  # make sure optimizer is initialized
        sd = optimizer.state_dict()
        load_checkpoint(sd)  # copy state directly into tensors, re-shard if needed
        optimizer.load_state_dict(sd)  # replace param_groups
docstr_tok: ["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",".","Because","of","introduced","strictness","it","allows","us","to",":","*","do","compatibility","checks","for","state","and","param_groups",",","which","improves","usability","*","avoid","state","duplication","by","directly","copying","into","state","tensors",",","e.g",".","optimizer.step","(",")","#","make","sure","optimizer","is","initialized","sd","=","optimizer.state_dict","(",")","load_checkpoint","(","sd",")","#","copy","state","directly","into","tensors",",","re-shard","if","needed","optimizer.load_state_dict","(","sd",")","#","replace","param_groups"]
code: def load_state_dict(self, state_dict: Mapping[str, Any]) -> None:
      """
      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.
  
      Because of introduced strictness it allows us to:
          * do compatibility checks for state and param_groups, which improves usability
          * avoid state duplication by directly copying into state tensors, e.g.
            optimizer.step()  # make sure optimizer is initialized
            sd = optimizer.state_dict()
            load_checkpoint(sd)  # copy state directly into tensors, re-shard if needed
            optimizer.load_state_dict(sd)  # replace param_groups
      """
  
      new_state = state_dict["state"]
      state = self.state
      params = self.params
  
      # Load state
      if len(state) != len(new_state):
          raise ValueError(
              f"Different parameter count: {len(state)} vs {len(new_state)}"
          )
      for param_key, param in params.items():
          if param not in state:
              continue
          if param_key not in new_state:
              raise ValueError(f"Parameter {param_key} not found")
          if len(state[param]) != len(new_state[param_key]):
              raise ValueError(
                  f"Different state size: {len(state[param])} vs {len(new_state[param_key])}"
              )
  
          KeyedOptimizer._update_param_state_dict_object(
              current_param_state_dict=state[param],
              param_state_dict_to_load=new_state[param_key],
              parent_keys=[param_key],
          )
  
      # Load param_groups.
      if self.defaults["_save_param_groups"]:
          new_param_groups = state_dict["param_groups"]
          param_groups = self.param_groups
  
          if len(param_groups) != len(new_param_groups):
              raise ValueError(
                  f"Different param_groups count: {len(param_groups)} vs {len(new_param_groups)}"
              )
          param_to_key = {param: key for key, param in params.items()}
          group_map = {}
          for group in param_groups:
              param_keys = []
              for param in group["params"]:
                  param_keys.append(param_to_key[param])
              group_map["/".join(sorted(param_keys))] = group
          new_group_map = {}
          for new_group in new_param_groups:
              param_keys = []
              for param_key in new_group["params"]:
                  param_keys.append(param_key)
              new_group_map["/".join(sorted(param_keys))] = new_group
          for group_key, group in group_map.items():
              if group_key not in new_group_map:
                  raise ValueError(f"Group {group_key} not found")
              new_group = new_group_map[group_key]
              if len(group) != len(new_group):
                  raise ValueError(
                      f"Different param_group size: {len(group)} vs {len(new_group)}"
                  )
              for k in group:
                  if k not in new_group:
                      raise ValueError(
                          f"Group key {k} not found for group {group_key}"
                      )
                  if k != "params":
                      group[k] = deepcopy(new_group[k])
  
      self.post_load_state_dict()
  
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 6 for calling the KeyedOptimizer obj's underlying member method code with example usage: obj.load_state_dict(state_dict) without return types