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_batch_prioritized_routing(self):
      set_seed(0)
      config = NllbMoeConfig(
          num_experts=4,
          hidden_size=32,
          d_ff=16,
          expert_capacity=4,
          second_expert_policy="random",
      )
      mask = torch.zeros(
          (self.batch_size * self.sequence_length), dtype=torch.bool
      )
      logits = torch.rand((self.batch_size * self.sequence_length, 4))
      config.batch_prioritized_routing = True
      router = NllbMoeTop2Router(config)
      top_1_mask, _ = router.route_tokens(logits, padding_mask=mask)
      # check that the routing is batch first. One of the last token is routed while expert capacity is very small
      # this means that it had a greater probability of being routed
      assert top_1_mask[-1, 0] == 1
  
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 2 for calling the NllbMoeRouterTest obj's underlying member method code with example usage: obj.test_batch_prioritized_routing() without return types