class_docstr: ``_NamedOptimizer`` takes a dict of parameters and exposes ``state_dict`` by parameter key.
  
  We replace the original key (number) in an optim to the
  fully qualified name (FQN) string. User can initialize the optim as they
  initialize a PyTorch optim, the only difference is that they also need to
  pass in the FQN of each parameters.
  
  Args:
      named_parameters (Mapping[str, Union[torch.Tensor, ShardedTensor]]):
          Mapping from FQN to parameter.
      optimizer_class (optim.Optimizer):
          The class of optimizer to instantiate.
      param_groups (Collection[Mapping[str, Any]]):
          `param_groups` to pass to optimizer if specified.
          The key of the inner map needs to be FQNs.
          Default: None
      module (nn.Module): the module whose parameters to updated
          by the optimizer.
      args: arguments to pass to the optimizer constructor.
      kwargs: arguments to pass to the optimizer constructor.
  
  Example::
      >>> # xdoctest: +SKIP("distributed")
      >>> from torch import optim
      >>> from torch.distributed.optim import _NamedOptimizer
      >>>
      >>> # Define the named optimizer.
      >>> m = Model(...)
      >>> named_optim = _NamedOptimizer(m.named_parameters(), optim.SGD)
      >>> # Forward pass + backward pass.
      >>> named_optim.step()
      >>> ...
      >>> # Call state_dict for the named optimizer returns a FQN state_dict.
      >>> named_optim.state_dict()
  
  Warning: This API is still in development and subject to change.
  
  TODO: Add tutorial for _NamedOptimizer.
  TODO: Add documentation in the docstring for the public attributes
        like self.param_groups and self.named_parameters.
class_docstr_tok: ["``","_NamedOptimizer","``","takes","a","dict","of","parameters","and","exposes","``","state_dict","``","by","parameter","key",".","We","replace","the","original","key","(","number",")","in","an","optim","to","the","fully","qualified","name","(","FQN",")","string",".","User","can","initialize","the","optim","as","they","initialize","a","PyTorch","optim",",","the","only","difference","is","that","they","also","need","to","pass","in","the","FQN","of","each","parameters",".","Args",":","named_parameters","(","Mapping","[","str",",","Union","[","torch.Tensor",",","ShardedTensor","]","]",")",":","Mapping","from","FQN","to","parameter",".","optimizer_class","(","optim.Optimizer",")",":","The","class","of","optimizer","to","instantiate",".","param_groups","(","Collection","[","Mapping","[","str",",","Any","]","]",")",":","`","param_groups","`","to","pass","to","optimizer","if","specified",".","The","key","of","the","inner","map","needs","to","be","FQNs",".","Default",":","None","module","(","nn.Module",")",":","the","module","whose","parameters","to","updated","by","the","optimizer",".","args",":","arguments","to","pass","to","the","optimizer","constructor",".","kwargs",":","arguments","to","pass","to","the","optimizer","constructor",".","Example",":",":",">",">",">","#","xdoctest",":","+SKIP","(","``","distributed","''",")",">",">",">","from","torch","import","optim",">",">",">","from","torch.distributed.optim","import","_NamedOptimizer",">",">",">",">",">",">","#","Define","the","named","optimizer",".",">",">",">","m","=","Model","(","...",")",">",">",">","named_optim","=","_NamedOptimizer","(","m.named_parameters","(",")",",","optim.SGD",")",">",">",">","#","Forward","pass","+","backward","pass",".",">",">",">","named_optim.step","(",")",">",">",">","...",">",">",">","#","Call","state_dict","for","the","named","optimizer","returns","a","FQN","state_dict",".",">",">",">","named_optim.state_dict","(",")","Warning",":","This","API","is","still","in","development","and","subject","to","change",".","TODO",":","Add","tutorial","for","_NamedOptimizer",".","TODO",":","Add","documentation","in","the","docstring","for","the","public","attributes","like","self.param_groups","and","self.named_parameters","."]
docstr: Return the ``state_dict`` of the optimizer.
  
  Instead of using number to index
  parameters, we will use module fully qualified name (FQN) as the key.
docstr_tok: ["Return","the","``","state_dict","``","of","the","optimizer",".","Instead","of","using","number","to","index","parameters",",","we","will","use","module","fully","qualified","name","(","FQN",")","as","the","key","."]
code: def state_dict(self) -> Dict[str, Any]:
      """
      Return the ``state_dict`` of the optimizer.
  
      Instead of using number to index
      parameters, we will use module fully qualified name (FQN) as the key.
      """
      state_dict = self._optimizer.state_dict()
      param_groups = state_dict["param_groups"]
  
      ret_state = {
          self.ordered_param_keys[st_key]: state_val
          for st_key, state_val in state_dict["state"].items()
      }
  
      ret_groups = []
      for group in param_groups:
          param_keys = []
          for param in group["params"]:
              param_keys.append(self.ordered_param_keys[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)
  
      return self._post_state_dict(
          {"state": ret_state, "param_groups": ret_groups}
      )
  
imports: import logging
  import warnings
  from copy import deepcopy
  from typing import Any, Callable, Collection, Dict, List, Mapping, Optional, Union, overload
  import torch
  import torch.nn
  from torch import optim
  from torch.distributed._shard.sharded_tensor import ShardedTensor
  from torch.distributed.fsdp import FullyShardedDataParallel
label_desc: Use image node_id 3 for calling the _NamedOptimizer obj's underlying member method code with example usage: obj.state_dict() and returns: self