veloxquant_mlx.quantizers.polarquant module

PolarQuant — recursive polar-coordinate KV quantizer.

Applies a random rotation followed by RecursivePolarTransform’s recursive angle/radius decomposition (see transforms/polar.py), then quantizes each level’s angles with a dedicated per-level codebook and stores the final scalar radius directly. Reconstruction and inner-product estimation both go through explicit decode of the rotated vector, unlike TurboQuant’s residual-correction quantizers.

class veloxquant_mlx.quantizers.polarquant.PolarQuantizer(d, b=2, m=128, seed=42, n_levels=4, store=None, use_hadamard=False, **kwargs)

Bases: Quantizer

PolarQuant: recursive polar-coordinate quantization.

Pipeline (encode):

y = Π·x (random rotation) angles, r_final = polar_transform(y) For each level ℓ: angle_idx[ℓ] = codebook[ℓ].quantize(angles[ℓ]) Store: angle_idx (list), r_final (fp16 scalar per token)

Pipeline (decode):

For each level ℓ: angles[ℓ] = codebook[ℓ].dequantize(angle_idx[ℓ]) y_hat = polar_inverse(angles, r_final) x_hat = Π^T · y_hat

Parameters:

• d (int) – Input dimension. Must be divisible by 2^n_levels.

• b (int) – Bit-width for each level.

• n_levels (int) – Number of polar recursion levels (default 4).

• seed (int) – Random seed.

• m (int) – Unused (for API consistency).

• store (ArtifactStore | None) – Optional ArtifactStore.

• use_hadamard (bool) – If True, use randomized Hadamard instead of QR rotation.

• kwargs (Any)

encode(x)

Encode via rotation, polar transform, and per-level codebook lookup.

Parameters:

x (Any) – Array of shape (batch, d), fp16.

Return type:

EncodedVector

Returns:

EncodedVector with angles (list of index arrays) and final_radius.

decode(ev)

Reconstruct approximate vector from polar encoding.

Parameters:

ev (EncodedVector) – EncodedVector with angles (indices) and final_radius.

Return type:

Any

Returns:

Reconstructed array of shape (batch, d), fp16.

estimate_inner_product(q, ev)

Estimate ⟨q, k⟩ by decoding then computing dot product.

Parameters:

• q (Any) – Query vector, shape (d,) or (1, d).

• ev (EncodedVector) – Encoded keys.

Return type:

Any

Returns:

Estimated inner products, shape (batch,), fp16.