docstr: Check that NamedOptimizer exposes the expected state dict
  interface when multiple param groups are specified.
docstr_tok: ["Check","that","NamedOptimizer","exposes","the","expected","state","dict","interface","when","multiple","param","groups","are","specified","."]
code: def test_state_dict_multi_param_group(self):
      """Check that NamedOptimizer exposes the expected state dict
      interface when multiple param groups are specified."""
      m = TestDummyModel()
      m_dup = TestDummyModel()
      optim_1 = torch.optim.SGD(
          [
              {"params": m.net1.parameters()},
              {"params": m.net3.parameters(), "lr": 1e-3},
          ],
          lr=1e-2,
          momentum=0.9,
      )
  
      optim_2 = torch.optim.Adam(
          [
              {"params": m.net2.parameters()},
              {"params": m.net4.parameters(), "lr": 1e-5},
          ]
      )
  
      named_optim_1 = _NamedOptimizer(
          m_dup.named_parameters(),
          torch.optim.SGD,
          [
              {"params": m_dup.net1.parameters()},
              {"params": m_dup.net3.parameters(), "lr": 1e-3},
          ],
          lr=1e-2,
          momentum=0.9,
      )
  
      named_optim_2 = _NamedOptimizer(
          m_dup.named_parameters(),
          torch.optim.Adam,
          [
              {"params": m_dup.net2.parameters()},
              {"params": m_dup.net4.parameters(), "lr": 1e-5},
          ],
      )
      self._compare_param_groups(
          optim_1.param_groups, named_optim_1.param_groups
      )
      self._compare_param_groups(
          optim_2.param_groups, named_optim_2.param_groups
      )
  
      _run_model_training(
          [
              (m, [optim_1, optim_2]),
              (m_dup, [named_optim_1, named_optim_2]),
          ]
      )
      self._compare_param_groups(
          optim_1.param_groups, named_optim_1.param_groups
      )
      self._compare_param_groups(
          optim_2.param_groups, named_optim_2.param_groups
      )
      sd_1 = optim_1.state_dict()
      sd_2 = optim_2.state_dict()
      named_sd_1 = named_optim_1.state_dict()
      named_sd_2 = named_optim_2.state_dict()
  
      # Compare "state" in optim state dict
      self._compare_state_dict_group(
          sd_1["state"][0],
          named_sd_1["state"]["net1.0.weight"],
          assert_equal=True,
      )
      self._compare_state_dict_group(
          sd_2["state"][1],
          named_sd_2["state"]["net2.0.bias"],
          assert_equal=True,
      )
      self._compare_state_dict_group(
          sd_1["state"][2],
          named_sd_1["state"]["net3.weight"],
          assert_equal=True,
      )
      self._compare_state_dict_group(
          sd_2["state"][3],
          named_sd_2["state"]["net4.1.bias"],
          assert_equal=True,
      )
  
      # Compare "param_groups" in optim state dict
      self._compare_state_dict_group(
          sd_1["param_groups"][0],
          named_sd_1["param_groups"][0],
          assert_equal=True,
      )
      self._compare_state_dict_group(
          sd_2["param_groups"][1],
          named_sd_2["param_groups"][1],
          assert_equal=True,
      )
  
imports: import unittest
  import torch
  import torch.nn
  from torch.distributed.optim import _NamedOptimizer
label_desc: Use image node_id 5 for calling the NamedOptimizerTest obj's underlying member method code with example usage: obj.test_state_dict_multi_param_group() without return types