code: def test_router_forward(
      router: MoeRouter,
      batch_size: int,
      seq_len: int,
      num_experts: int,
      num_groups: int,
  ):
      x = torch.randn((batch_size * seq_len, num_experts)).cuda()
      if num_groups > 1:
          x = x.expand(num_groups, -1, -1)
  
      router.train()
      if isinstance(router, TopKRouter):
          _, combine_array, dispatch_mask = router(x, expert_capacity=2)
      else:
          _, combine_array, dispatch_mask = router(x)
      assert combine_array.shape[:-1] == x.shape
      assert dispatch_mask.shape[:-1] == x.shape
      assert torch.all(dispatch_mask.sum(-1).sum(-1) <= router.k_value)
  
      router.eval()
      if isinstance(router, TopKRouter):
          _, combine_array, dispatch_mask = router(x, expert_capacity=2)
      else:
          _, combine_array, dispatch_mask = router(x)
      assert combine_array.shape[:-1] == x.shape
      assert dispatch_mask.shape[:-1] == x.shape
      assert torch.all(dispatch_mask.sum(-1).sum(-1) <= router.k_value)
  
label_desc: Use image node_id 1 for calling a global function with example usage: test_router_forward(router, batch_size, seq_len, num_experts, num_groups) without return types