(self, d: int, bits: int)
| 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) |
nothing calls this directly
no test coverage detected