Method quantize

turboquant_cutile/codebook.py:70–72

(self, x: torch.Tensor)

        self.centroids, self.boundaries = solve_lloyd_max(d, bits)

    def quantize(self, x: torch.Tensor) -> torch.Tensor:
        diffs = x.unsqueeze(-1) - self.centroids.to(x.device)
        return diffs.abs().argmin(dim=-1)

    def dequantize(self, indices: torch.Tensor) -> torch.Tensor:
        return self.centroids.to(indices.device)[indices.long()]