code: def test_constructor(self):
      assert isinstance(self.switch_mlp_smm, SwitchMLP)
      assert isinstance(self.switch_mlp_gmm, GroupedMLP)
  
      num_weights_smm = sum(
          [p.numel() for p in self.switch_mlp_smm.parameters()]
      )
      num_weights_gmm = sum(
          [p.numel() for p in self.switch_mlp_gmm.parameters()]
      )
  
      # For the same hyper-parm model configs except the `moe_grouped_gemm`,
      # GroupedGEMM and sequential GEMMs should hold the same number of parms.
      assert num_weights_smm == num_weights_gmm
      # expected num weights: router linear weights+bias + MLP weights(no bias) of all experts
      expected_num_weights = (
          self.hidden_size * self.num_experts
          + self.hidden_size
          * (self.fc1_ffn_hidden_size + self.fc2_ffn_hidden_size)
          * self.num_experts
      )
      assert num_weights_smm == expected_num_weights
  
      assert torch.equal(
          self.switch_mlp_smm.router.weight,
          self.switch_mlp_gmm.router.weight,
      )
  
      # weight1: [h, num_experts*4h]
      # weight2: [num_experts*4h, h]
      assert self.switch_mlp_gmm.weight1.shape[0] == self.hidden_size
      assert (
          self.switch_mlp_gmm.weight1.shape[1]
          == self.num_experts * self.fc1_ffn_hidden_size
      )
      if self.gated_linear_unit:
          assert (
              self.switch_mlp_gmm.weight2.shape[0]
              == self.num_experts * self.fc2_ffn_hidden_size
          )
          assert (
              self.switch_mlp_gmm.weight2.shape[1] == self.hidden_size
          )
      else:
          assert (
              self.switch_mlp_gmm.weight1.shape
              == self.switch_mlp_gmm.weight2.t().shape
          )
  
file_path: megatron-lm/tests/unit_tests/transformer/moe/test_grouped_mlp.py
imports: import pytest
  import torch
  import torch.nn.functional
  from megatron.arguments import parse_args
  from megatron.core.models.gpt.gpt_layer_specs import get_gpt_layer_with_transformer_engine_spec
  from megatron.core.transformer.moe.grouped_mlp import GroupedMLP
  from megatron.core.transformer.moe.switch_mlp import SwitchMLP
  from megatron.core.transformer.transformer_config import TransformerConfig
  from megatron.initialize import _set_random_seed
  from megatron.model import Float16Module
  from tests.unit_tests.test_utilities import Utils
label_desc: Use image node_id 3 for calling the TestParallelGroupedMLP obj's underlying member method code with example usage: obj.test_constructor() without return types