class torchelie.nn.MultiVQ(latent_dim: int, num_tokens: int, num_codebooks: int, dim: int = 1, commitment: float = 0.25, init_mode: str = 'normal', return_indices: bool = True, max_age: int = 1000)

Multi codebooks quantization layer from Neural Discrete Representation Learning

  • latent_dim (int) – number of features along which to quantize

  • num_tokens (int) – number of tokens in the codebook

  • num_codebooks (int) – number of parallel codebooks

  • dim (int) – dimension along which to quantize an angular distance

  • return_indices (bool) – whether to return the indices of the quantized code points

forward(x: torch.Tensor)Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]
training: bool