mixle.models.unified_quantizer module

Unified per-tensor quantization surface with a method picker (roadmap I1).

Unifies TWO existing, independently-landed quantization mechanisms behind one interface:

  1. mixle.task.quantize – int8/int4 per-tensor symmetric quantization (quantize_dequantize_array(), the exact core of quantize_mlp) and LNS (log-number-system) compute quantization (LogNumberSystem). Both are “already in-tree and load-bearing” (roadmap context); this module wraps them, it does not reimplement their arithmetic.

  2. mixle.models.sorted_profile_quantizer (roadmap G4) – head-exact + parametric-tail per-tensor storage, honestly scoped to optimizer-states / KV-tails / anomaly-detection, NOT a general weight quantizer (see that module’s docstring).

quantize_tensor() is the single per-tensor entry point: explicit method= values ("int8", "int4", "lns", "sorted_profile") dispatch directly to the corresponding underlying primitive; method="auto" runs a small picker – reusing this codebase’s existing UCB1 discrete-arm machinery (the D5/ConditionalJIT “small learned/bandit controller picks an action per context” pattern, replicated here at per-tensor scale: the “context” is one tensor, the “arms” are the four methods, the “reward” is measured reconstruction quality at a matched byte budget) – to choose, PER TENSOR, whichever method gives the best reconstruction at the requested size budget.

Every QuantizedTensor – explicit or auto-picked – carries a QuantizationReceipt: the chosen method, its measured bytes/error/compression ratio, and (for auto-pick) the SAME real numbers for every method that was considered and rejected, so no choice is silently unexplained.

Matched-size protocol. All four methods are compared at a shared byte budget target_bytes = ceil(n * bits / 8) (bits defaults to 8, i.e. the int8 rate). int8/int4 hit their fixed rate by construction (8 or 4 bits/element); LNS quantizes log-magnitude to the same integer width (nibble-packed at 4 bits, reusing mixle.task.quantize._pack_nibbles()) plus a packed sign bit per element (mixle.models.unified_quantizer does not compress the sign, so LNS carries a small, honestly-reported n/8-byte overhead on top of the magnitude bits); sorted-profile has no size KNOB tied to bits at all – its rate is set by the tensor’s own permutation-index dtype (mixle.models.sorted_profile_quantizer._index_dtype()), so at some tensor sizes it will not fit the budget at all. A method whose ACTUAL measured nbytes exceeds target_bytes is marked eligible=False in the receipt and is never auto-picked, even if its reconstruction error would otherwise be the best – “matched size” is enforced on real measured bytes, not assumed ones.

class QuantizationReceipt(method, auto, nbytes, reconstruction_error, compression_ratio, target_bytes, candidates=<factory>, notes='')[source]

Bases: object

Explains why method was used for one tensor: its own measured numbers, plus – for auto-pick – the same real numbers for every OTHER method that was considered and rejected.

Parameters:

  • method (str)

  • auto (bool)

  • nbytes (int)

  • reconstruction_error (float)

  • compression_ratio (float)

  • target_bytes (int)

  • candidates (dict[str, MethodCandidate])

  • notes (str)

rejected()[source]

The candidates NOT chosen (empty for an explicit, non-auto dispatch).

Return type:

dict[str, MethodCandidate]

class MethodCandidate(method, nbytes, reconstruction_error, compression_ratio, eligible, reward)[source]

Bases: object

Real, measured numbers for one method considered for one tensor – the raw material every QuantizationReceipt (chosen or rejected) is built from.

Parameters:
class QuantizedTensor(method, shape, payload, receipt)[source]

Bases: object

Unified result: whichever underlying encoding was produced, with a .reconstruct() that works regardless of which method was actually used, plus the receipt explaining the choice.

Parameters:

  • method (str)

  • shape (tuple[int, ...])

  • payload (SymmetricQuantPayload | LNSTensorPayload | SortedProfileEncoding)

  • receipt (QuantizationReceipt)

class SymmetricQuantPayload(wq, scale, bits)[source]

Bases: object

The int8/int4 payload: exactly what mixle.task.quantize.quantize_dequantize_array() returns.

Parameters:
class LNSTensorPayload(codes, sign_bits, step, center, bits, n)[source]

Bases: object

The LNS payload: log-magnitude quantized via mixle.engines.lns.LogNumberSystem (the SAME class mixle.task.quantize.lns_classifier uses for its integer log-space inference), plus a packed sign bit per element (LNS strips sign when it takes log|v|).

codes stores the quantized log-magnitude: a raw int8 array at bits=8, or – reusing mixle.task.quantize._pack_nibbles(), the SAME nibble packer QuantizedMLP uses for its int4 weights – two-per-byte packed at bits=4, so LNS’s on-disk rate genuinely matches int4’s, not just its accounting.

Parameters:
quantize_tensor(tensor, method='auto', *, bits=8, target_compression=None, top_k=0, tail_family=None, gof_threshold=DEFAULT_GOF_THRESHOLD, clip_percentile=None, seed=None)[source]

The single per-tensor quantization entry point (roadmap I1).

Parameters:

  • tensor (Any) – numpy array or torch tensor of any shape.

  • method (str) – "auto" (default, runs the picker) or one of METHODS ("int8", "int4", "lns", "sorted_profile") to dispatch directly to the corresponding underlying primitive.

  • bits (int) – target bits/element for the matched-size budget (target_bytes = ceil(n*bits/8)); also the bit width int8/int4/LNS quantize AT when explicitly requested. target_compression (if given) overrides bits as bits = 32 // target_compression (e.g. target_compression=4 -> 8 bits, matching the fp32 -> int8 4x-compression convention).

  • top_k (int) – forwarded to mixle.models.sorted_profile_quantizer.fit_sorted_profile().

  • tail_family (Any) – forwarded to mixle.models.sorted_profile_quantizer.fit_sorted_profile().

  • gof_threshold (float) – forwarded to mixle.models.sorted_profile_quantizer.fit_sorted_profile().

  • clip_percentile (float | None) – forwarded to mixle.task.quantize.quantize_dequantize_array() for the int8/int4 methods.

  • seed (int | None) – seed for the auto-pick UCB1 picker (deterministic either way, kept for interface symmetry with the rest of mixle’s bandit call sites).

  • target_compression (int | None)

Returns:

QuantizedTensor

Return type:

QuantizedTensor

lns_quantize_array(flat, bits=8)[source]

Quantize a flat float array in the log-magnitude domain via LogNumberSystem.

mixle.engines.lns.LogNumberSystem was built to quantize log-DENSITIES (already-log-domain, already-positive-support values); a general tensor has both sign and a linear-domain magnitude. This function is the honest generalization: split v = sign(v) * |v|, quantize log(|v| + eps) with LogNumberSystem (the exact same quantize/dequantize integer machinery lns_classifier uses – not reimplemented here), and pack the sign separately (1 bit/element via numpy.packbits()). This is the natural fit for multiplicative-scale / heavy-tailed data (LNS’s whole reason to exist), and a poor fit for already-near-zero-centered, additive-scale data – exactly where int8’s LINEAR quantization should (and, per the model-zoo test, does) win instead.

Parameters:
Return type:

LNSTensorPayload

lns_dequantize_array(payload)[source]

Inverse of lns_quantize_array().

Parameters:

payload (LNSTensorPayload)

Return type:

ndarray