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","."]
code: def add_param_group(self, param_group: Any) -> None:
      raise NotImplementedError()
  
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 7 for calling the KeyedOptimizer obj's underlying member method code with example usage: obj.add_param_group(param_group) without return types