class_docstr: Switch Transformers has different blocks from classic transformer based models.
  The Swift MLP contains a Router class, that has to be tested to check if it is correctly implemented
  
  Original implementation of the routers here:
class_docstr_tok: ["Switch","Transformers","has","different","blocks","from","classic","transformer","based","models",".","The","Swift","MLP","contains","a","Router","class",",","that","has","to","be","tested","to","check","if","it","is","correctly","implemented","Original","implementation","of","the","routers","here",":"]
code: def test_top_2_routing(self):
      # test routing with minimal reproduction
      mask = torch.ones(
          (self.batch_size, self.sequence_length), dtype=torch.bool
      )
      mask[0][0] = False
      mask[1][0] = False
      mask = mask.reshape(-1)
      set_seed(0)
      hidden_states = torch.rand(
          (
              self.batch_size,
              self.sequence_length,
              self.config.hidden_size,
          )
      )
      classfier = torch.nn.Linear(
          self.config.hidden_size, self.config.num_experts
      )
      hf_router = NllbMoeTop2Router(self.config)
  
      _, _, hidden_dim = hidden_states.shape
      logits = classfier(
          hidden_states.reshape(
              (self.batch_size * self.sequence_length), hidden_dim
          )
      )
      top_1_mask, router_probs = hf_router.route_tokens(
          logits, padding_mask=mask
      )
      torch.argmax(top_1_mask, dim=-1)
      router_mask = router_probs.bool()
      set_seed(0)
      experts = [
          torch.nn.Linear(hidden_dim, hidden_dim),
          torch.nn.Linear(hidden_dim, hidden_dim),
          torch.nn.Linear(hidden_dim, hidden_dim),
          torch.nn.Linear(hidden_dim, hidden_dim),
      ]
      hidden_states = hidden_states.reshape(
          (self.batch_size * self.sequence_length), hidden_dim
      )
      masked_hidden_states = torch.einsum(
          "bm,be->ebm", hidden_states, router_mask
      )
      for idx, expert in enumerate(experts):
          token_indices = router_mask[:, idx]
          combining_weights = router_probs[token_indices, idx]
          expert_output = expert(
              masked_hidden_states[idx, token_indices]
          )
          expert_output *= 1 - self.config.moe_token_dropout
          masked_hidden_states[idx, token_indices] = torch.einsum(
              "b,be->be", combining_weights, expert_output
          )
      hidden_states = masked_hidden_states.sum(dim=0).reshape(
          self.batch_size, self.sequence_length, hidden_dim
      )
  
      EXPECTED_MEAN_FAIRSEQ_HIDDEN_STATES = torch.Tensor([[ 7.0340e-04,  2.7997e-03, -1.3351e-02, -7.6705e-03, -3.5089e-03,3.9773e-03,  7.4593e-03,  1.2566e-02,  3.5860e-03, -2.7448e-02,-1.3731e-02, -1.0534e-02, -1.3606e-02, -1.5048e-02, -2.8914e-03,-5.0371e-03, -1.3963e-03,  6.0076e-03, -1.1380e-02, -1.4620e-02, 5.2401e-03,  8.4660e-04, -1.5319e-03, -1.6735e-02,  1.1302e-02, 3.6119e-03,  4.6084e-03, -1.3458e-02,  7.7792e-05,  1.4312e-02, 4.9107e-03, -5.0936e-03], [-4.4538e-03,  3.1026e-03,  1.4121e-04, -4.8121e-03, -5.6279e-03, 7.2493e-03,  3.9769e-03,  1.1114e-02, -1.5666e-03, -2.3477e-02, 8.7268e-03,  1.3446e-02, -2.8845e-05, -1.7287e-02,  8.7619e-03, -4.5316e-03, -1.2164e-02,  5.7461e-03, -4.5861e-03, -9.3907e-03, 2.9808e-02,  8.9206e-04, -7.6232e-04, -1.4173e-02,  3.0208e-03, 1.5310e-02,  9.7717e-03,  3.1014e-03,  7.8042e-03,  8.0197e-03, 3.4784e-03, -7.1728e-03]])  # fmt: skip
      self.assertTrue(
          torch.allclose(
              hidden_states.mean(1),
              EXPECTED_MEAN_FAIRSEQ_HIDDEN_STATES,
              1e-4,
          )
      )
  
imports: import copy
  import tempfile
  import unittest
  from transformers import NllbMoeConfig, is_torch_available, set_seed
  from transformers.testing_utils import require_sentencepiece, require_tokenizers, require_torch, require_torch_fp16, slow, torch_device
  from transformers.utils import cached_property
  from ...generation.test_utils import GenerationTesterMixin
  from ...test_configuration_common import ConfigTester
  from ...test_modeling_common import ModelTesterMixin, ids_tensor
  from ...test_pipeline_mixin import PipelineTesterMixin
label_desc: Use image node_id 1 for calling the NllbMoeRouterTest obj's underlying member method code with example usage: obj.test_top_2_routing() without return types