Method encode

model/kronos.py:142–159

Encodes the input data into quantized indices. Args: x (torch.Tensor): Input tensor of shape (batch_size, seq_len, d_in). half (bool, optional): Whether to use half quantization in BSQuantizer. Defaults to False. Returns: torch.Tensor: Q

(self, x, half=False)

        return x

    def encode(self, x, half=False):
        """
        Encodes the input data into quantized indices.

        Args:
            x (torch.Tensor): Input tensor of shape (batch_size, seq_len, d_in).
            half (bool, optional): Whether to use half quantization in BSQuantizer. Defaults to False.

        Returns:
            torch.Tensor: Quantized indices from BSQuantizer.
        """
        z = self.embed(x)
        for layer in self.encoder:
            z = layer(z)
        z = self.quant_embed(z)

        bsq_loss, quantized, z_indices = self.tokenizer(z, half=half, collect_metrics=False)
        return z_indices

    def decode(self, x, half=False):
        """