Function solve_lloyd_max

turboquant_cutile/codebook.py:20–58

Returns (centroids, boundaries) as sorted float32 tensors. centroids: (2^bits,) boundaries: (2^bits - 1,)

(
    d: int,
    bits: int,
    max_iter: int = 200,
    tol: float = 1e-10,
)

def solve_lloyd_max(
    d: int,
    bits: int,
    max_iter: int = 200,
    tol: float = 1e-10,
) -> tuple[torch.Tensor, torch.Tensor]:
    """
    Returns (centroids, boundaries) as sorted float32 tensors.
    centroids: (2^bits,)  boundaries: (2^bits - 1,)
    """
    n_levels = 1 << bits
    sigma = 1.0 / math.sqrt(d)
    pdf = lambda x: _gaussian_pdf(x, sigma)

    lo, hi = -3.5 * sigma, 3.5 * sigma
    centroids = [lo + (hi - lo) * (i + 0.5) / n_levels for i in range(n_levels)]

    for _ in range(max_iter):
        boundaries = [
            (centroids[i] + centroids[i + 1]) / 2.0 for i in range(n_levels - 1)
        ]
        edges = [lo * 3] + boundaries + [hi * 3]
        new_centroids = []
        for i in range(n_levels):
            a, b = edges[i], edges[i + 1]
            num, _ = integrate.quad(lambda x: x * pdf(x), a, b)
            den, _ = integrate.quad(pdf, a, b)
            new_centroids.append(num / den if den > 1e-15 else centroids[i])
        if max(abs(new_centroids[i] - centroids[i]) for i in range(n_levels)) < tol:
            break
        centroids = new_centroids

    boundaries = [
        (centroids[i] + centroids[i + 1]) / 2.0 for i in range(n_levels - 1)
    ]
    return (
        torch.tensor(centroids, dtype=torch.float32),
        torch.tensor(boundaries, dtype=torch.float32),
    )


class LloydMaxCodebook: