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","."]
docstr: Unflatten params. If the current instance is already unflattened, then
  it will remain unflattened after the context manager exits.
  
  Args:
      flat_params (List[Tensor], Optional):
          flat params to use for unflattening.
          If provided, the current instance must be in a flattened state
          at the start of the context manager. The provided Tensor must be
          appropriately sized and will only be used within the context
          manager. After the context manager exits, we will revert to
          using ``self.flat_params``
          Default: None.
docstr_tok: ["Unflatten","params",".","If","the","current","instance","is","already","unflattened",",","then","it","will","remain","unflattened","after","the","context","manager","exits",".","Args",":","flat_params","(","List","[","Tensor","]",",","Optional",")",":","flat","params","to","use","for","unflattening",".","If","provided",",","the","current","instance","must","be","in","a","flattened","state","at","the","start","of","the","context","manager",".","The","provided","Tensor","must","be","appropriately","sized","and","will","only","be","used","within","the","context","manager",".","After","the","context","manager","exits",",","we","will","revert","to","using","``","self.flat_params","``","Default",":","None","."]
code: def unflatten_params(
      self, flat_params: Optional[List[Tensor]] = None
  ) -> Generator:
      """
      Unflatten params. If the current instance is already unflattened, then
      it will remain unflattened after the context manager exits.
  
      Args:
          flat_params (List[Tensor], Optional):
              flat params to use for unflattening.
              If provided, the current instance must be in a flattened state
              at the start of the context manager. The provided Tensor must be
              appropriately sized and will only be used within the context
              manager. After the context manager exits, we will revert to
              using ``self.flat_params``
              Default: None.
      """
      assert (
          flat_params is None or self.is_flattened
      ), "Unflattening with external flat_param requires current instance to be flattened"
  
      orig_flattened = self.is_flattened
      if orig_flattened:
          orig_flat_params = self.flat_params
          self._unflatten_params(
              cast(Optional[List[Optional[Tensor]]], flat_params)
          )
  
      # Put yield in a try...finally in case the caller catches the exception and handles
      # it. In that case, we need to properly handle the undoing of state here.
      try:
          yield
      finally:
          if orig_flattened:
              self._flatten_params(orig_flat_params)
  
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 10 for calling the FlattenParamsWrapper obj's underlying member method code with example usage: obj.unflatten_params(flat_params) without return types