def test_equivalency_token_chose_masked_router(self):
        r"""
        This test tests the equivalency between the `SwitchTransformersTop1Router`
        originally implemented from here: TODO: provide link
        """

        input_tokens = torch.Tensor(
            [
                [
                    [0.6433916, 0.18188512, 0.02240455, 0.563781],
                    [0.5526401, 0.0958724, 0.34253013, 0.03644359],
                    [0.08744538, 0.7909105, 0.35205448, 0.53364205],
                ],
                [
                    [0.02900076, 0.4168595, 0.5802449, 0.91486526],
                    [0.27414513, 0.14991808, 0.9383501, 0.5209162],
                    [0.51207185, 0.90618336, 0.7309413, 0.95533276],
                ],
            ]
        )

        model = SwitchTransformersTop1Router(self.config)

        model.classifier.weight = torch.nn.Parameter(
            torch.Tensor(
                [
                    [0.02008116, 0.00620062],
                    [-0.00811031, -0.00031623],
                    [-0.03542127, 0.02703803],
                    [0.02335377, -0.02971946],
                ],
            ).t()
        )

        expert_index, _, router_logits = model(input_tokens)
        router_probs = torch.softmax(router_logits, dim=-1)

        router_z_loss = router_z_loss_func(router_logits)
        auxiliary_loss = load_balancing_loss_func(router_probs, torch.argmax(expert_index, dim=-1))

        self.assertAlmostEqual(auxiliary_loss.item(), 1.000308, places=5)
        self.assertAlmostEqual(router_z_loss.item(), 0.4789799, places=5)

        # self.assertTrue(torch.allclose(expert_index.bool().unsqueeze(-1), expected_dispatch_mask))

        This test tests the equivalency between the `SwitchTransformersTop1Router`
        originally implemented from here: TODO: provide link
        
dpkd/transformers/tests/models/switch_transformers/test_modeling_switch_transformers.py