class_docstr: Takes a dict of parameters and exposes state_dict by parameter key.
  
  Convenience wrapper to take in optim_factory callable to create KeyedOptimizer
class_docstr_tok: ["Takes","a","dict","of","parameters","and","exposes","state_dict","by","parameter","key",".","Convenience","wrapper","to","take","in","optim_factory","callable","to","create","KeyedOptimizer"]
code: def __init__(
      self,
      params: Mapping[str, Union[torch.Tensor, ShardedTensor]],
      optim_factory: OptimizerFactory,
  ) -> None:
      self._optimizer: optim.Optimizer = optim_factory(
          list(params.values())
      )
      super().__init__(
          params, self._optimizer.state, self._optimizer.param_groups
      )
  
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 1 to create a new KeyedOptimizerWrapper object from inherited base classes: KeyedOptimizer with example: obj = KeyedOptimizerWrapper(params, optim_factory)