class_docstr: A wrapper for transparently flattening a Module's parameters.
  
  Compared to the original implementation [1], this version:
  - removes tracing
  - supports shared parameters
  - handles state_dict/load_state_dict transparently
  - is renamed to FlattenParamsWrapper
  - refactored to use the FlatParameter class
  - extended to support flattening multiple groups of params (useful
    when different groups of params need different hyperparameters, like
    learning rate or weight decay)
  
  [1] https://github.com/SsnL/PyTorch-Reparam-Module
  
  Args:
      module (nn.Module):
          The module to wrap.
      param_list (Optional[List[List[nn.Parameter]]]):
          Only flatten parameters appearing in the given groups.
          If the param_list is an empty list, then no parameters will get flattened.
          Note, if a single param is in one of the list, it still get flattened and the
          original param is removed and replaced with the flatten one.
          Default: None, flatten all parameters (if any)
      flat_param_names (Optional[List[str]]):
          originally, give each flat_param a unique name. Note a "flat_param_"
          prefix will be added to those names.
class_docstr_tok: ["A","wrapper","for","transparently","flattening","a","Module","'s","parameters",".","Compared","to","the","original","implementation","[","1","]",",","this","version",":","-","removes","tracing","-","supports","shared","parameters","-","handles","state_dict\/load_state_dict","transparently","-","is","renamed","to","FlattenParamsWrapper","-","refactored","to","use","the","FlatParameter","class","-","extended","to","support","flattening","multiple","groups","of","params","(","useful","when","different","groups","of","params","need","different","hyperparameters",",","like","learning","rate","or","weight","decay",")","[","1","]","https",":","\/\/github.com\/SsnL\/PyTorch-Reparam-Module","Args",":","module","(","nn.Module",")",":","The","module","to","wrap",".","param_list","(","Optional","[","List","[","List","[","nn.Parameter","]","]","]",")",":","Only","flatten","parameters","appearing","in","the","given","groups",".","If","the","param_list","is","an","empty","list",",","then","no","parameters","will","get","flattened",".","Note",",","if","a","single","param","is","in","one","of","the","list",",","it","still","get","flattened","and","the","original","param","is","removed","and","replaced","with","the","flatten","one",".","Default",":","None",",","flatten","all","parameters","(","if","any",")","flat_param_names","(","Optional","[","List","[","str","]","]",")",":","originally",",","give","each","flat_param","a","unique","name",".","Note","a","``","flat_param_","''","prefix","will","be","added","to","those","names","."]
code: def __init__(
      self,
      module: nn.Module,
      param_list: ParamGroups = None,
      flat_param_names: Optional[List[str]] = None,
  ):
      super().__init__()
      self._fpw_module = module
      self.is_flattened = False
  
      # Handle param_list being None.
      if param_list is None:
          param_list = list(module.parameters())
  
      # Be backward compatible and turn a single param list into a list of
      # a single list.
      if len(param_list) > 0 and isinstance(
          param_list[0], nn.Parameter
      ):
          param_list = [cast(List[nn.Parameter], param_list)]
  
      # Since the parameters will be deleted, let's record the number original
      # parameters managed by this class. This and get_param_views function
      # below are used by fsdp_optim_utils.py to save/restore optimizer state,
      # which mirrors the flatten parameters here.
      self.num_params_managed = 0
  
      self._param_sets = []
      overall_param_set: Set[nn.Parameter] = set()
      for p_list in param_list:
          # Remove any duplicates from the list.
          p_set: Set[nn.Parameter] = set(
              cast(List[nn.Parameter], p_list)
          )
  
          self.num_params_managed += len(p_set)
          overall_param_set = overall_param_set.union(p_set)
  
          # Convert from list of Parameters to set of (Module, name) tuples,
          # which will survive in case the parameter instances are reset.
          # Also, a shared param will correctly appear under multiple modules
          # as they should.
          new_p_set_with_names = set()
          for m in self.modules():
              for n, p in m.named_parameters(recurse=False):
                  if p in p_set:
                      new_p_set_with_names.add((m, n))
          if new_p_set_with_names:
              self._param_sets.append(new_p_set_with_names)
  
      if len(overall_param_set) != self.num_params_managed:
          # Each p_list above could have shared params. However, you can't
          # have shared params cross different p_list. That means part of
          # the flattened parameter must be shared, which is impossible to
          # support.
          raise ValueError(
              f"Incorrect param groups {len(overall_param_set)} vs {self.num_param_managed}"
          )
  
      self.flat_params: List[nn.Parameter] = []
  
      # Prepare flat param names.
      if flat_param_names is None:
          flat_param_names = [
              f"{i}" for i, _ in enumerate(self._param_sets)
          ]
      if len(flat_param_names) != len(self._param_sets):
          raise ValueError(
              "Names and number of param lists must be equal"
          )
      if len(flat_param_names) != len(set(flat_param_names)):
          raise ValueError(
              "Each flat param must be given a unique name"
          )
      self.flat_param_names = [
          f"flat_param_{n}" for n in flat_param_names
      ]
      flat_param: Optional[nn.Parameter] = None
  
      # Init all flat_params.
      for new_p_set in self._param_sets:
          (
              params,
              param_infos,
              shared_param_infos,
          ) = self._init_flatten_params(new_p_set)
          flat_param = FlatParameter(params, params[0].requires_grad)
          flat_param._param_infos = param_infos
          flat_param._shared_param_infos = shared_param_infos
          self.flat_params.append(flat_param)
  
      self._flatten_params(self.flat_params)
  
      # Register hook to be called after state_dict() to remove the
      # "_fpw_module." prefix and before load_state_dict() to add it back.
      self._register_state_dict_hook(_post_state_dict_hook)
      self._register_load_state_dict_pre_hook(_pre_load_state_dict_hook)
  
      # Flag to indicate whether state_dict() should automatically unflatten
      # params. This defaults to True, but may be set to False if the user
      # explicitly requests a flat state dict via flat_state_dict().
      self._auto_unflatten_state_dict = True
  
imports: from contextlib import contextmanager
  from itertools import chain
  import typing
  from typing import TYPE_CHECKING, Any, Dict, Generator, Iterator, List, Mapping, NamedTuple, Optional, Sequence, Set, Tuple, Union, cast
  import torch
  from torch import Tensor
  import torch.nn
  from fairscale.internal.state_dict import replace_by_prefix_
label_desc: Use image node_id 1 to create a new FlattenParamsWrapper object from inherited base classes: nn with example: obj = FlattenParamsWrapper(module, param_list, flat_param_names)