(
self,
dim: int,
codebook_size: int,
codebook_dim: tp.Optional[int] = None,
decay: float = 0.99,
epsilon: float = 1e-5,
kmeans_init: bool = True,
kmeans_iters: int = 50,
threshold_ema_dead_code: int = 2,
commitment_weight: float = 1.0,
)
| 249 | """ |
| 250 | |
| 251 | def __init__( |
| 252 | self, |
| 253 | dim: int, |
| 254 | codebook_size: int, |
| 255 | codebook_dim: tp.Optional[int] = None, |
| 256 | decay: float = 0.99, |
| 257 | epsilon: float = 1e-5, |
| 258 | kmeans_init: bool = True, |
| 259 | kmeans_iters: int = 50, |
| 260 | threshold_ema_dead_code: int = 2, |
| 261 | commitment_weight: float = 1.0, |
| 262 | ): |
| 263 | super().__init__() |
| 264 | _codebook_dim: int = default(codebook_dim, dim) |
| 265 | |
| 266 | requires_projection = _codebook_dim != dim |
| 267 | self.project_in = ( |
| 268 | nn.Linear(dim, _codebook_dim) if requires_projection else nn.Identity() |
| 269 | ) |
| 270 | self.project_out = ( |
| 271 | nn.Linear(_codebook_dim, dim) if requires_projection else nn.Identity() |
| 272 | ) |
| 273 | |
| 274 | self.epsilon = epsilon |
| 275 | self.commitment_weight = commitment_weight |
| 276 | |
| 277 | self._codebook = EuclideanCodebook( |
| 278 | dim=_codebook_dim, |
| 279 | codebook_size=codebook_size, |
| 280 | kmeans_init=kmeans_init, |
| 281 | kmeans_iters=kmeans_iters, |
| 282 | decay=decay, |
| 283 | epsilon=epsilon, |
| 284 | threshold_ema_dead_code=threshold_ema_dead_code, |
| 285 | ) |
| 286 | self.codebook_size = codebook_size |
| 287 | |
| 288 | @property |
| 289 | def codebook(self): |
nothing calls this directly
no test coverage detected