veloxquant_mlx.codebooks package

Scalar codebooks and their centroid-fitting strategies for quantizers.

Groups ScalarCodebook (nearest-centroid quantize/dequantize), CodebookFactory (builds a codebook for a given distribution and bit-width), and the CodebookStrategy implementations that fit centroids via Lloyd-Max on a Gaussian, Beta, or polar-angle PDF, or space them uniformly — the pieces every scalar-quantization method in quantizers/ and cache/ composes to turn continuous coordinates into b-bit codes.

class veloxquant_mlx.codebooks.CodebookFactory[source]

Bases: object

Factory for creating ScalarCodebook instances via registered strategies.

All instantiation of codebooks should go through this factory.

Example:

cb = CodebookFactory.create("gaussian", b=2, d=128)
static create(distribution, b, d, polar_level=1)[source]

Create a ScalarCodebook for the given distribution and bit-width.

Parameters:

  • distribution (Literal['gaussian', 'beta', 'polar_level', 'uniform']) – Centroid distribution type. - "gaussian": N(0, 1/d) — high-d TurboQuant. - "beta": exact Beta distribution — low-d TurboQuant. - "polar_level": polar angle distribution (requires polar_level). - "uniform": uniformly spaced centroids.

  • b (int) – Bit-width (1 <= b <= 8).

  • d (int) – Vector dimension.

  • polar_level (int) – Polar recursion level (only used when distribution=”polar_level”).

Return type:

ScalarCodebook

Returns:

Configured ScalarCodebook instance, memoized by (distribution, b, d, polar_level) (VeloxQuant-MLX#508) — centroid construction is a pure, deterministic function of these args, and ScalarCodebook is read-only after construction, so it’s safe to share across callers.

Raises:

QuantizerConfigError – If parameters are invalid.

class veloxquant_mlx.codebooks.ScalarCodebook(centroids)[source]

Bases: Codebook

MLX-based scalar codebook for nearest-centroid quantisation.

Quantisation uses broadcast argmin over the codebook (O(d·k) per batch). Dequantisation is a simple gather (O(d) per batch).

For large k (> VoronoiTree.LINEAR_THRESHOLD), a VoronoiTree provides O(log k) per-coordinate search in pure Python (used only when MLX is unavailable or for verification).

Parameters:

centroids (ndarray) – 1-D array of sorted centroid values, shape (k,).

property k: int

Number of centroids (2^b).

property b: int

Bit-width of this codebook.

quantize(y)[source]

Map coordinates to nearest-centroid indices via boundary search.

dist_{i,j,k} = |y_{i,j} - c_k| → argmin_k

Parameters:

y (Any) – Input array of shape (batch, d), fp16.

Return type:

Any

Returns:

Index array of shape (batch, d), dtype uint8.

dequantize(idx)[source]

Retrieve centroid values via gather.

Parameters:

idx (Any) – Index array of shape (batch, d), dtype uint8.

Return type:

Any

Returns:

Centroid values of shape (batch, d), fp16.

nearest_numpy(value)[source]

O(log k) nearest-centroid lookup for scalar queries.

Falls back to linear for k <= 16. Used for verification.

Parameters:

value (float) – Scalar float query.

Return type:

int

Returns:

Index of the nearest centroid.

centroids_numpy()[source]

Return centroid array as float32 numpy array.

Return type:

ndarray

centroids_mx()[source]

Return centroid array as MLX fp16 array.

Return type:

Any

boundaries_mx()[source]

Return sorted Voronoi boundaries (midpoints) as MLX fp16 array.

Shape (k - 1,) — the same array quantize() compares against. Exposed for callers (e.g. fused Metal kernels) that need to replicate quantize’s boundary-count semantics outside MLX.

Return type:

Any

class veloxquant_mlx.codebooks.LloydMaxGaussianStrategy(support_sigma_factor=6.0)[source]

Bases: CodebookStrategy

Optimal codebook for N(0, 1/d) — high-d TurboQuant coordinates.

Uses Lloyd-Max algorithm with a Gaussian PDF scaled to sigma = 1/sqrt(d).

Parameters:

support_sigma_factor (float) – Support range as multiples of sigma (default 5).

compute_centroids(b, d)[source]

Compute 2^b centroids for N(0, 1/d) distribution.

Parameters:

  • b (int) – Bit-width.

  • d (int) – Vector dimension.

Return type:

ndarray

Returns:

Sorted centroid array of shape (2^b,), float64.

class veloxquant_mlx.codebooks.LloydMaxBetaStrategy[source]

Bases: CodebookStrategy

Optimal codebook for the exact Beta-distributed coordinates of TurboQuant.

Uses the analytically correct Beta distribution for the given dimension d. Preferred for d < 64 where the Gaussian approximation is inaccurate.

Parameters:

None.

compute_centroids(b, d)[source]

Compute 2^b centroids for Beta(d/2, d/2) coordinate distribution.

Parameters:

  • b (int) – Bit-width.

  • d (int) – Vector dimension.

Return type:

ndarray

Returns:

Sorted centroid array of shape (2^b,), float64.

class veloxquant_mlx.codebooks.PolarAngleSamplingStrategy(level=1)[source]

Bases: CodebookStrategy

Optimal codebook for PolarQuant angle distribution at each level.

Uses Lloyd-Max on the analytically derived polar angle PDF. The level must be set via the level attribute before calling compute_centroids().

Parameters:

level (int) – Polar recursion level (1-indexed, must be >= 1).

compute_centroids(b, d)[source]

Compute 2^b centroids for polar angle PDF at the configured level.

Parameters:

  • b (int) – Bit-width.

  • d (int) – Vector dimension (unused; distribution is dimension-independent).

Return type:

ndarray

Returns:

Sorted centroid array of shape (2^b,), float64.

class veloxquant_mlx.codebooks.UniformStrategy(lo=-1.0, hi=1.0)[source]

Bases: CodebookStrategy

Uniform (midpoint) codebook — baseline for comparison.

Divides the support range into 2^b equal-width cells and places centroids at cell midpoints.

Parameters:

  • lo (float) – Lower bound of the uniform support (default -1.0).

  • hi (float) – Upper bound of the uniform support (default 1.0).

compute_centroids(b, d)[source]

Compute 2^b uniformly spaced centroids.

Parameters:

  • b (int) – Bit-width.

  • d (int) – Unused.

Return type:

ndarray

Returns:

Sorted centroid array of shape (2^b,), float64.

Submodules