(self)
| 24 | assert out.shape == (1,) |
| 25 | |
| 26 | def test_gradient_flow(self): |
| 27 | router = TokenRouter(hidden_dim=128, bottleneck_dim=32) |
| 28 | x = torch.randn(4, 128, requires_grad=True) |
| 29 | out = router(x) |
| 30 | loss = out.sum() |
| 31 | loss.backward() |
| 32 | assert x.grad is not None |
| 33 | assert x.grad.shape == (4, 128) |
| 34 | |
| 35 | def test_different_bottleneck_dims(self): |
| 36 | for bd in [32, 64, 128, 256]: |
nothing calls this directly
no test coverage detected