Class TokenRouter

python/TIDE/router.py:11–19

class TokenRouter(nn.Module):
    def __init__(self, hidden_dim: int, bottleneck_dim: int = 128):
        super().__init__()
        self.down = nn.Linear(hidden_dim, bottleneck_dim, bias=False)
        self.up = nn.Linear(bottleneck_dim, 1, bias=False)
        self.act = nn.SiLU()

    def forward(self, hidden_state: torch.Tensor) -> torch.Tensor:
        return torch.sigmoid(self.up(self.act(self.down(hidden_state)))).squeeze(-1)


@dataclass