mixle.represent.quantize module

VectorQuantizer – learn a discrete vocabulary IN the shared embedding space, don’t guess it upstream.

Discrete tokens, when you want them (compression, transfer, a fixed vocabulary), come after embedding, not before segmentation: fit a codebook to the continuous vectors and each vector’s nearest code is its token id. The codebook is a learned model (k-means / a mixture), so the vocabulary is inferred from data rather than assumed – and because every modality is embedded into the same space, one codebook is a cross-modal vocabulary (an image patch and a word can share a token id when they land near the same centroid).

fit/quantize/dequantize are the codec; straight_through gives the VQ-VAE gradient so the codebook and the encoders can be trained end to end under a generative or downstream objective. This is the only place discreteness lives – the segmenter and embedding stay vocabulary-free.

class VectorQuantizer(num_codes, dim, *, seed=0)

Bases: object

A learned codebook over R^dim: nearest-centroid quantization of embedding vectors into discrete ids.

Parameters:

• num_codes (int)

• dim (int)

• seed (int)

property codebook: ndarray | None

The fitted (num_codes, dim) centroids – read-only, or None before fit().

The codebook IS the learned vocabulary: every token id this codec emits or decodes is defined by it, so it is fitted state, not a scratch buffer. It used to be a plain public attribute, and a single post-fit write silently redefined the vocabulary (MXR-080-1906). Verified: after vq.fit(x), setting vq.codebook[0, 0] = nan collapsed quantize from [0, 2, 0, 0, 2] to [0, 0, 0, 0, 0] – every vector assigned to code 0, because a NaN distance never wins an argmin – and turned reconstruction_error into nan rather than an error. Rebinding was just as open: vq.codebook = np.zeros((99, 7)) was accepted beside an unchanged dim=2 and num_codes=3.

The array is frozen with writeable = False and reachable only through this property, which matches mixle.engines.formats.CodebookFormat, whose codebook is copied and sealed the same way. dequantize returns codebook[ids], and NumPy fancy indexing copies, so decoded vectors are still ordinary writable arrays.

fit(vectors, *, iters=25)

Fit the codebook by k-means (Lloyd) on vectors (n, dim) – the vocabulary is learned, not assumed.

num_codes is a cap, not a guarantee: fitting on fewer than num_codes samples yields one center per sample, and num_codes is updated in place to that actual count so it – and every bounds check derived from it (dequantize, a caller building a one-hot over the vocabulary) – stays honest about the codebook’s real capacity instead of quietly overstating it.

iters must be a positive integer: iters=0 used to publish the random initialization centers as a fitted codec, with no Lloyd step ever run.

Parameters:

• vectors (ndarray)

• iters (int)

Return type:

VectorQuantizer

quantize(vectors)

Nearest-code id for each vector – the discrete token stream (n,).

Parameters:

vectors (ndarray)

Return type:

ndarray

dequantize(ids)

Codebook vectors for token ids (n,) -> (n, dim) (the reconstruction / de-tokenization).

Ids must already be integral. The cast to int64 used to happen before the range check, so a fractional id such as 0.9 truncated to the ordinary in-range token 0 and decoded as if it had been produced by quantize().

Parameters:

ids (ndarray)

Return type:

ndarray

reconstruction_error(vectors)

Mean squared quantization error – the codebook’s fidelity (a codebook-size / bitrate knob).

Parameters:

vectors (ndarray)

Return type:

float

straight_through(vectors)

VQ-VAE straight-through estimator: return quantized vectors but pass gradients to vectors unchanged.

Lets the encoders and (with a codebook-commitment loss) the codebook train end to end through the discrete bottleneck. vectors is a torch tensor (n, dim).

Parameters:

vectors (Any)

Return type:

Any