Source code for veloxquant_mlx.quantizers.composite

"""Mixed-precision quantizer that splits a vector's channels by magnitude.

``CompositeQuantizer`` implements outlier-aware quantization: a small set
of high-magnitude "outlier" channels is routed to a high-bit-width child
quantizer while the remaining "inlier" channels go to a low-bit-width
child, and the two child ``EncodedVector``s are nested under one composite
encoding. This lets any two existing per-channel quantizers be combined to
protect the few dimensions that dominate reconstruction error without
paying full bit-width across the whole vector.
"""

from __future__ import annotations

from typing import Any

import mlx.core as mx
import numpy as np

from veloxquant_mlx.core.abstractions import Quantizer
from veloxquant_mlx.core.context import EncodedVector




class CompositeQuantizer(Quantizer):
    """Split-channel quantizer that routes outlier and inlier channels separately.

    Outlier channels are quantized with a high-bit-width quantizer.
    Inlier channels are quantized with a low-bit-width quantizer.

    The outlier_encoded and inlier_encoded fields of the returned EncodedVector
    nest the children's outputs.

    Args:
        outlier_quantizer: Quantizer for high-magnitude (outlier) channels.
        inlier_quantizer: Quantizer for low-magnitude (inlier) channels.
        outlier_idx: Integer indices of outlier channels in the original vector.
        total_dim: Total dimension d of the original vector.
    """

    def __init__(
        self,
        outlier_quantizer: Quantizer,
        inlier_quantizer: Quantizer,
        outlier_idx: np.ndarray,
        total_dim: int,
    ) -> None:
        self._outlier_q = outlier_quantizer
        self._inlier_q = inlier_quantizer
        self._outlier_idx = np.asarray(outlier_idx, dtype=np.int32)
        self._inlier_idx = np.setdiff1d(np.arange(total_dim), self._outlier_idx)
        self._d = total_dim
        # MLX-array copies of the same indices, precomputed once so decode()'s
        # scatter (mx array indexed by mx array) never converts NumPy -> MLX
        # on the read hot path.
        self._outlier_idx_mx = mx.array(self._outlier_idx)
        self._inlier_idx_mx = mx.array(self._inlier_idx.astype(np.int32))



    def encode(self, x: Any) -> EncodedVector:
        """Encode by routing channels to respective child quantizers.

        Args:
            x: Array of shape (batch, d), fp16.

        Returns:
            CompositeEncodedVector with nested outlier and inlier encodings.
        """
        if x.ndim == 1:
            x = x[None]

        x_out = x[:, self._outlier_idx_mx]
        x_in = x[:, self._inlier_idx_mx]

        ev_out = self._outlier_q.encode(x_out)
        ev_in = self._inlier_q.encode(x_in)

        return EncodedVector(
            quantizer_type="composite",
            batch_size=x.shape[0],
            dim=self._d,
            outlier_idx=self._outlier_idx,
            outlier_encoded=ev_out,
            inlier_encoded=ev_in,
        )




    def decode(self, ev: EncodedVector) -> Any:
        """Decode and reassemble channels.

        Args:
            ev: CompositeEncodedVector.

        Returns:
            Reconstructed array of shape (batch, d), fp16.
        """
        if ev.outlier_encoded is None or ev.inlier_encoded is None:
            raise ValueError(
                "CompositeQuantizer.decode: ev.outlier_encoded/inlier_encoded is "
                "None — ev wasn't produced by this quantizer's encode()."
            )

        x_out = self._outlier_q.decode(ev.outlier_encoded)  # (batch, n_out)
        x_in = self._inlier_q.decode(ev.inlier_encoded)  # (batch, n_in)

        batch = x_out.shape[0]
        out = mx.zeros((batch, self._d), dtype=x_out.dtype)
        out[:, self._outlier_idx_mx] = x_out
        out[:, self._inlier_idx_mx] = x_in.astype(x_out.dtype)
        return out




    def estimate_inner_product(self, q: Any, ev: EncodedVector) -> Any:
        """Estimate ⟨q, k⟩ by summing contributions from both children.

        Args:
            q: Query vector, shape (d,) or (1, d).
            ev: CompositeEncodedVector.

        Returns:
            Estimated inner products, shape (batch,), fp16.
        """
        if ev.outlier_encoded is None or ev.inlier_encoded is None:
            raise ValueError(
                "CompositeQuantizer.estimate_inner_product: "
                "ev.outlier_encoded/inlier_encoded is None — ev wasn't produced "
                "by this quantizer's encode()."
            )

        q_flat = q.reshape(-1)
        q_out = q_flat[self._outlier_idx_mx]
        q_in = q_flat[self._inlier_idx_mx]

        ip_out = self._outlier_q.estimate_inner_product(q_out, ev.outlier_encoded)
        ip_in = self._inlier_q.estimate_inner_product(q_in, ev.inlier_encoded)
        return ip_out + ip_in


    def __repr__(self) -> str:
        return (
            f"CompositeQuantizer("
            f"n_outliers={len(self._outlier_idx)}, "
            f"n_inliers={len(self._inlier_idx)}, "
            f"d={self._d})"
        )