Pre-solved Lloyd-Max codebook for a given (d, bits) pair.
| 59 | |
| 60 | |
| 61 | class LloydMaxCodebook: |
| 62 | """Pre-solved Lloyd-Max codebook for a given (d, bits) pair.""" |
| 63 | |
| 64 | def __init__(self, d: int, bits: int): |
| 65 | self.d = d |
| 66 | self.bits = bits |
| 67 | self.n_levels = 1 << bits |
| 68 | self.centroids, self.boundaries = solve_lloyd_max(d, bits) |
| 69 | |
| 70 | def quantize(self, x: torch.Tensor) -> torch.Tensor: |
| 71 | diffs = x.unsqueeze(-1) - self.centroids.to(x.device) |
| 72 | return diffs.abs().argmin(dim=-1) |
| 73 | |
| 74 | def dequantize(self, indices: torch.Tensor) -> torch.Tensor: |
| 75 | return self.centroids.to(indices.device)[indices.long()] |
| 76 | |
| 77 | def __repr__(self) -> str: |
| 78 | return ( |
| 79 | f"LloydMaxCodebook(d={self.d}, bits={self.bits}, " |
| 80 | f"levels={self.n_levels})" |
| 81 | ) |
no outgoing calls