class_docstr: Base test class and used for CPU case.
docstr: Testing some parameters are flatten, with others left non-flatten.
docstr_tok: ["Testing","some","parameters","are","flatten",",","with","others","left","non-flatten","."]
code: def test_partial_flattening(self):
      """Testing some parameters are flatten, with others left non-flatten."""
      module = self._get_transformer()
      num_params = sum(p.numel() for p in module.parameters())
  
      params_to_flatten = list(
          module.encoder.layers[1].parameters()
      ) + list(module.decoder.layers[0].parameters())
      num_params_to_flatten = sum(p.numel() for p in params_to_flatten)
  
      module = FlattenParamsWrapper(
          module, param_list=params_to_flatten
      )
      assert module.flat_param.numel() == num_params_to_flatten
      assert sum(p.numel() for p in module.parameters()) == num_params
  
      # flattened parameters are removed
      assert len(list(module.encoder.layers[1].parameters())) == 0
      assert len(list(module.decoder.layers[0].parameters())) == 0
  
      # non-flattened parameters remain
      assert len(list(module.encoder.layers[0].parameters())) > 0
      assert len(list(module.decoder.layers[1].parameters())) > 0
  
      # test that changing the module dtype works properly
      orig_dtype = params_to_flatten[0].dtype
      new_dtype = (
          torch.float32
          if orig_dtype == torch.float16
          else torch.float16
      )
      assert module.flat_param.dtype == orig_dtype
      assert all(
          p.dtype == orig_dtype
          for p in module.encoder.layers[0].parameters()
      )
      module = module.to(dtype=new_dtype)
      assert module.flat_param.dtype == new_dtype
      assert all(
          p.dtype == new_dtype
          for p in module.encoder.layers[0].parameters()
      )
  
imports: from collections import OrderedDict
  import unittest
  import torch
  from fairscale.fair_dev.testing.testing import objects_are_equal
  from fairscale.nn import FlattenParamsWrapper
label_desc: Use image node_id 12 for calling the TestFlattenParams obj's underlying member method code with example usage: obj.test_partial_flattening() without return types