VQ

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

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

• dim (int) – dimension along which to quantize

• 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]]

Forward pass

Parameters

x (tensor) – input tensor

Returns

quantized tensor, or (quantized tensor, indices) if self.return_indices

resample_dead(x)

update_usage(indices)

commitment: float

dim: int

embedding: torch.nn.modules.sparse.Embedding

init_mode: str

initialized: torch.Tensor

return_indices: bool