MultiVQ

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

Parameters

• 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