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: Build metadata for need-to-be-flatten parameters and returns a list
      contains the need-to-be-flatten parameters.
  
      This also returns param_infos and shared_param_infos, which
      will be attached to the flat parameter object.
  
  Args:
      p_set (set):
          A set of (module, param_name) for a set of params that needed
          to be flattened. There could be shared params in this set.
docstr_tok: ["Build","metadata","for","need-to-be-flatten","parameters","and","returns","a","list","contains","the","need-to-be-flatten","parameters",".","This","also","returns","param_infos","and","shared_param_infos",",","which","will","be","attached","to","the","flat","parameter","object",".","Args",":","p_set","(","set",")",":","A","set","of","(","module",",","param_name",")","for","a","set","of","params","that","needed","to","be","flattened",".","There","could","be","shared","params","in","this","set","."]
code: def _init_flatten_params(
      self, p_set: Set[Tuple[nn.Module, str]]
  ) -> Tuple[
      List[nn.Parameter],
      List[Tuple[str, nn.Module, str]],
      List[Tuple[str, str, nn.Module, str, nn.Module, str]],
  ]:
      """Build metadata for need-to-be-flatten parameters and returns a list
          contains the need-to-be-flatten parameters.
  
          This also returns param_infos and shared_param_infos, which
          will be attached to the flat parameter object.
  
      Args:
          p_set (set):
              A set of (module, param_name) for a set of params that needed
              to be flattened. There could be shared params in this set.
      """
      param_infos = []
      shared_param_memo: Dict[
          nn.Parameter, Tuple[str, nn.Module, str]
      ] = {}
      shared_param_infos = []
      params = []
      fp32 = []
      fp16 = []
      for module_name, m in self.named_modules():
          for n, p in m.named_parameters(recurse=False):
              if p.dtype != torch.float16:
                  fp32.append(module_name)
              else:
                  fp16.append(module_name)
              if p is not None and (m, n) in p_set:
                  if p in shared_param_memo:
                      mname, shared_m, shared_n = shared_param_memo[p]
                      shared_param_infos.append(
                          (module_name, mname, m, n, shared_m, shared_n)
                      )
                  else:
                      shared_param_memo[p] = (module_name, m, n)
                      param_infos.append((module_name, m, n))
                      params.append(p)
      del shared_param_memo
      fp16_msg, fp32_msg = ",".join(fp16), ",".join(fp32)
      assert (
          len(set(p.dtype for p in params)) == 1
      ), f"expects all parameters to have same dtype: fp32: {fp32_msg} \n fp16: {fp16_msg} "
      assert (
          len(set(p.requires_grad for p in params)) == 1
      ), f"expects all parameters to have same requires_grad {p_set}"
      assert len(params) == len(
          set(params)
      ), "params list should not have dups"
      return params, param_infos, shared_param_infos
  
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 4 for calling the FlattenParamsWrapper obj's underlying member method code with example usage: obj._init_flatten_params(p_set) and returns: params,  param_infos,  shared_param_infos