Struct ResidualCodec 

pub struct ResidualCodec {
    pub nbits: u32,
    pub dim: usize,
    pub centroids: Vec<f32>,
    pub bucket_cutoffs: Vec<f32>,
    pub bucket_weights: Vec<f32>,
}

A trained residual-quantization codec.

Cutoffs partition the real line into 2^nbits buckets. bucket_cutoffs holds the 2^nbits - 1 internal boundaries in ascending order; bucket_weights holds the 2^nbits reconstruction values used when decoding. Both are codec-wide: the same cutoffs/weights are applied to every residual dimension of every token.

Fields

nbits: u32

Number of bits per residual dimension. Typically 2 or 4.

dim: usize

Dimensionality of the (original) token embeddings.

centroids: Vec<f32>

Flat row-major coarse centroids, k × dim.

bucket_cutoffs: Vec<f32>

(2^nbits) - 1 ascending cutoff values for bucketing residuals.

bucket_weights: Vec<f32>

2^nbits reconstruction values, one per bucket.

Implementations

impl ResidualCodec

pub fn num_buckets(&self) -> usize

Number of buckets this codec partitions the residual space into.

pub fn num_centroids(&self) -> usize

Number of coarse centroids stored.

pub fn packed_bytes(&self) -> usize

Number of packed bytes each encoded vector uses.

pub fn validate(&self) -> Result<()>

Validate internal shape invariants. Called automatically by encode/decode; exposed so callers loading a codec from disk can fail fast.

Errors

Returns PlaidError::InvalidCodec with a description of the constraint that’s violated.

pub fn encode_vector(&self, vector: &[f32]) -> Result<EncodedVector>

Encode a single token embedding.

Finds the nearest centroid, computes the residual, and quantizes each dimension against bucket_cutoffs.

Errors

Returns PlaidError::InvalidCodec if this codec fails its shape invariants.

Panics

Panics if vector.len() != dim.

pub fn batch_encode_tokens(&self, tokens: &[f32]) -> Result<(Vec<u32>, Vec<u8>)>

Encode every token in a flat n × dim buffer in one batched pass, returning the per-token centroid id and a flat n × dim code buffer.

The expensive step — the nearest-centroid lookup — runs as a single matmul through crate::kmeans::assign_points, which uses candle’s GEMM (CPU or CUDA depending on build). The residual + bucket loop stays scalar because per-element searchsorted would otherwise require either a 3-D broadcast against the cutoffs table or a per-cutoff kernel launch — both less efficient than a tight Rust loop over the small cutoffs vector. Returning the codes flat avoids n Vec<u8> allocations; callers split into per-token slices as needed.

Errors

Returns PlaidError::InvalidCodec if the codec fails its shape invariants, or PlaidError::Tensor if the matmul-driven nearest-centroid lookup fails.

Panics

Panics if tokens.len() % dim != 0 or if tokens is empty.

pub fn batch_encode_tokens_with_chunk_rows( &self, tokens: &[f32], chunk_rows: usize, ) -> Result<(Vec<u32>, Vec<u8>)>

Same as [batch_encode_tokens] but processes the input in tiles of chunk_rows tokens per upload.

This is the path the PLAID builder uses for pools that would otherwise exceed VRAM — only [chunk_rows, dim] f32 ever lives on the device at once, so peak VRAM stays bounded in chunk_rows regardless of corpus size or embedding dimension. The output is byte-identical to [batch_encode_tokens] for any chunk_rows ≥ 1, which the prop_batch_encode_chunked_is_partition_invariant hegel property checks across shrunk-counterexample shapes.

Errors

Returns PlaidError::InvalidCodec on codec-shape violations, or PlaidError::Tensor if any per-chunk allocation or matmul fails.

Panics

Panics if tokens.len() % dim != 0 or if chunk_rows == 0.

pub fn batch_encode_tokens_on_tensor( &self, tokens_tensor: &Tensor, tokens: &[f32], ) -> Result<(Vec<u32>, Vec<u8>)>

Same as [batch_encode_tokens] but reuses a pre-uploaded tokens tensor for the nearest-centroid matmul.

The PLAID builder runs k-means on this same corpus immediately before calling batch encode; threading the device tensor through saves the second 3.47 GB host→device copy that would otherwise collide with the first one in cudarc’s caching allocator and OOM a 12 GB card.

tokens must be the host-side backing buffer for tokens_tensor — the residual + bit-pack loop is scalar and still reads token bytes from the host slice.

Errors

Returns PlaidError::InvalidCodec on codec-shape violations, or PlaidError::Tensor if the nearest-centroid matmul fails.

Panics

Panics if tokens.len() % dim != 0.

pub fn decode_vector(&self, encoded: &EncodedVector) -> Result<Vec<f32>>

Reconstruct an approximate token embedding from its codes.

Errors

Returns PlaidError::InvalidCodec if this codec fails its shape invariants.

Panics

Panics if codes.len() != dim or if any code is out of range.

pub fn decode_vector_with_table( &self, encoded: &EncodedVector, table: &DecodeTable, ) -> Result<Vec<f32>>

Decode an encoded vector using a pre-built DecodeTable.

Callers that decode many vectors in a row (e.g., the search path’s per-candidate decode loop) should build the table once outside the loop and reuse it here — each call then amounts to one table load per packed byte plus the centroid add.

Errors

Returns PlaidError::InvalidCodec if this codec fails its shape invariants.

pub fn reconstruction_error(&self, vector: &[f32]) -> Result<f32>

Return the squared L2 reconstruction error for vector under this codec. Useful as a lightweight codec-quality probe in tests and evaluation scripts.

Errors

Returns PlaidError::InvalidCodec if this codec fails its shape invariants.

Trait Implementations

impl Clone for ResidualCodec

fn clone(&self) -> Self

Returns a duplicate of the value.

fn clone_from(&mut self, source: &Self)

Performs copy-assignment from source.

impl Debug for ResidualCodec

fn fmt(&self, f: &mut Formatter<'_>) -> Result

Formats the value using the given formatter.