veloxquant_mlx.core.context module

Data-carrying types shared across the quantization handler chain and quantizers.

QuantizationContext is the mutable payload passed between QuantizationHandler stages (see veloxquant_mlx.handlers) as it is encoded or decoded; EncodedVector is the typed, memory-accountable output of Quantizer.encode(); TransformResult carries the per-level angles and radius produced by a recursive polar transform. These dataclasses decouple the handler pipeline and quantizer implementations from MLX’s array type via lazy import.

class veloxquant_mlx.core.context.QuantizationContext(x_original, mode, x_current, norm=None, rotated=None, indices=None, signs=None, residual_norm=None, angles=None, final_radius=None, outlier_idx=None, packed_bits=None, metadata=<factory>)[source]

Bases: object

Mutable payload flowing through a QuantizationHandler chain.

Variables:

  • x_original – Original input vectors, shape (batch, d).

  • mode – Whether the chain is encoding or decoding.

  • x_current – Working copy mutated by each handler.

  • norm – L2 norm stored by NormalizationHandler, shape (batch,).

  • rotated – Vector after RotationHandler, shape (batch, d).

  • indices – Codebook indices, shape (batch, d) uint8.

  • signs – QJL sign bits, shape (batch, m) int8.

  • residual_norm – Residual L2 norm for QJL stage, shape (batch,).

  • angles – Per-level polar angles from PolarTransformHandler.

  • final_radius – Scalar radius after all polar levels, shape (batch,).

  • outlier_idx – Channel positions of outlier channels.

  • packed_bits – Bit-packed index array from BitPackingHandler.

  • metadata – Arbitrary stage-specific metadata.

Parameters:

  • x_original (Any)

  • mode (Literal['encode', 'decode'])

  • x_current (Any)

  • norm (Any | None)

  • rotated (Any | None)

  • indices (Any | None)

  • signs (Any | None)

  • residual_norm (Any | None)

  • angles (list[Any] | None)

  • final_radius (Any | None)

  • outlier_idx (ndarray | None)

  • packed_bits (ndarray | None)

  • metadata (dict[str, Any])

x_original: Any
mode: Literal['encode', 'decode']
x_current: Any
norm: Any | None = None
rotated: Any | None = None
indices: Any | None = None
signs: Any | None = None
residual_norm: Any | None = None
angles: list[Any] | None = None
final_radius: Any | None = None
outlier_idx: ndarray | None = None
packed_bits: ndarray | None = None
metadata: dict[str, Any]
class veloxquant_mlx.core.context.EncodedVector(quantizer_type, batch_size, dim, indices=None, norm=None, signs=None, residual_norm=None, angles=None, final_radius=None, outlier_idx=None, outlier_encoded=None, inlier_encoded=None)[source]

Bases: object

Typed output of Quantizer.encode().

Different quantizer types populate different subsets of fields.

Variables:

  • quantizer_type – Registry key of the producing quantizer.

  • batch_size – Number of vectors encoded.

  • dim – Original vector dimensionality.

  • indices – Scalar codebook indices, shape (batch, d) uint8.

  • norm – Per-vector L2 norm, shape (batch,) fp16.

  • signs – QJL sign bits, shape (batch, m) int8.

  • residual_norm – QJL residual norm, shape (batch,) fp16.

  • angles – PolarQuant level angles, list of (batch, d/2^ℓ) arrays.

  • final_radius – PolarQuant scalar radius, shape (batch,) fp16.

  • outlier_idx – Outlier channel indices for CompositeQuantizer.

  • outlier_encoded – Nested EncodedVector for outlier channels.

  • inlier_encoded – Nested EncodedVector for inlier channels.

Parameters:
quantizer_type: str
batch_size: int
dim: int
indices: Any | None = None
norm: Any | None = None
signs: Any | None = None
residual_norm: Any | None = None
angles: list[Any] | None = None
final_radius: Any | None = None
outlier_idx: ndarray | None = None
outlier_encoded: EncodedVector | None = None
inlier_encoded: EncodedVector | None = None
memory_bytes()[source]

Compute exact memory footprint of this encoded representation.

Return type:

int

Returns:

Total bytes occupied by all stored arrays.

class veloxquant_mlx.core.context.TransformResult(angles, final_radius, n_levels)[source]

Bases: object

Output of RecursivePolarTransform.forward().

Variables:

  • angles – List of angle arrays, one per polar level. angles[0] has shape (batch, d/2), angles[ℓ] has shape (batch, d/2^(ℓ+1)).

  • final_radius – Scalar radius at the end of recursion, shape (batch,).

  • n_levels – Number of polar recursion levels applied.

Parameters:
angles: list[Any]
final_radius: Any
n_levels: int