pub trait MaxSim: StorageElement + CloneErrors
Error::KernelFailed when the kernel refuses to run, like on a buffer packed under
other Capabilities than the current ones.
MaxSim is the late-interaction primitive used by systems such as ColBERT. It is not “just another matrix multiply”.
use numkong::{MaxSimPackedMatrix, Tensor};
let queries = Tensor::<f32>::full(&[4, 16], 1.0).unwrap();
let docs = Tensor::<f32>::full(&[8, 16], 1.0).unwrap();
let queries_packed = MaxSimPackedMatrix::new(&queries).unwrap();
let docs_packed = MaxSimPackedMatrix::new(&docs).unwrap();
let score = queries_packed.score(&docs_packed).unwrap();
assert!(score.is_finite());API
TraitMaxSimTrait abstracting MaxSim pack/score operations per scalar type.5
ScoreScore type returned by MaxSim scoring.
type Score: Clone + Default;maxsim_pack_sizeReturns the packed buffer size in bytes for vectors vectors of given depth.
fn maxsim_pack_size(vectors: usize, depth: usize) -> Result<usize, Error>maxsim_packed_shapeReads the packed vector count and depth from the buffer header. # Safety packed must point to a buffer produced by maxsim_pack.
unsafe fn maxsim_packed_shape(packed: *const u8) -> Result<(usize, usize), Error>maxsim_packPack vectors into backend-specific quantized format.
unsafe fn maxsim_pack(
data: *const Self,
vectors: usize,
depth: usize,
stride: usize,
packed: *mut u8,
) -> Result<(), Error>Packing fuses no transposition, so depth-major vectors are transposed in a separate pass first.
Safety
maxsim_packedCompute MaxSim score on pre-packed buffers.
unsafe fn maxsim_packed(
queries: *const u8,
documents: *const u8,
query_count: usize,
document_count: usize,
depth: usize,
result: *mut Self::Score,
) -> Result<(), Error>Safety
- Both buffers must have been produced by
maxsim_packwith matching depth resultmust point to valid, writable memory forSelf::Score
StructMaxSimPackedMatrixPre-packed vector set for MaxSim scoring.16
pub struct MaxSimPackedMatrix<Scalar: MaxSim, Alloc: Allocator = Global>Both query and document vectors must be packed before scoring. The buffer uses i8 quantization for fast coarse screening, with full-precision originals retained for refinement.
Implements
Debug, Send, SynccloneClone this packed matrix, returning an error on allocation failure.
pub fn clone(&self) -> Result<Self, Error>empty_inAn empty packed set owning no allocation; fill it with pack_into.
pub fn empty_in(alloc: Alloc) -> Selfnew_inPack vectors from a 2D tensor view using a custom allocator.
pub fn new_in<Vectors, const MAX_RANK: usize>(data: &Vectors, alloc: Alloc) -> Result<Self, Error>
where
Vectors: TensorRef<Scalar, MAX_RANK> + ?Sized,Returns Err if the view is not 2D, the depth axis is not contiguous, the row stride is
negative, or allocation fails.
pack_intoRepack data into this set's existing buffer, reusing the allocation when the packed size fits capacity and reallocating through the stored allocator only when it must grow.
pub fn pack_into<Vectors, const MAX_RANK: usize>(&mut self, data: &Vectors) -> Result<(), Error>
where
Vectors: TensorRef<Scalar, MAX_RANK> + ?Sized,reservePre-grow the buffer to hold vectors vectors of depth, so a later pack_into that fits stays allocation-free with a stable pointer — hoist this out of a decode loop.
pub fn reserve(&mut self, vectors: usize, depth: usize) -> Result<(), Error>scoreCompute the MaxSim score — sum over queries of the max cosine to any document vector — against another packed matrix, treating self as the queries and other as the documents.
pub fn score<OtherAlloc: Allocator>(
&self,
other: &MaxSimPackedMatrix<Scalar, OtherAlloc>,
) -> Result<Scalar::Score, Error>Returns Err if:
- the two matrices were packed at different depths
- the kernel refuses a matrix, like one packed under other capabilities
allocatorReturns a reference to the allocator.
pub fn allocator(&self) -> &AllocvectorsNumber of vectors in the packed set.
pub fn vectors(&self) -> usizeshapeReturns the shape [vectors,depth] of the original vector set.
pub fn shape(&self) -> (usize, usize)pack_sizeBytes a packed buffer occupies for vectors vectors of the given depth under the active backend's layout, letting a caller pre-size an external buffer without packing.
pub fn pack_size(vectors: usize, depth: usize) -> Result<usize, Error>from_packed_bytes_inAdopt an externally-produced packed buffer by copying bytes into a container-owned allocation tagged with the given vectors and depth.
pub unsafe fn from_packed_bytes_in(
bytes: &[u8],
vectors: usize,
depth: usize,
alloc: Alloc,
) -> Result<Self, Error>The packed layout is not self-describing, so the caller must supply the shape the bytes were packed for.
Safety
capacityBytes currently allocated (>= the live packed size).
pub fn capacity(&self) -> usizeclearReset to logically empty, keeping the allocation so the next pack_into reuses it.
pub fn clear(&mut self)as_bytesReturns the packed data buffer.
pub fn as_bytes(&self) -> &[u8]as_ptrReturns a pointer to the packed data.
pub fn as_ptr(&self) -> *const u8newPack a 2D tensor of vectors for MaxSim scoring using the global allocator.
pub fn new<Vectors, const MAX_RANK: usize>(data: &Vectors) -> Result<Self, Error>
where
Vectors: TensorRef<Scalar, MAX_RANK> + ?Sized,The MaxSimPackedMatrix:: qualifier names the packing target — a tensor can be packed for
MaxSim or for dots, and those layouts differ, so construction goes through the typed
constructor rather than a bare tensor.pack().