diskann_quantization/multi_vector/
mod.rs

1/*
2 * Copyright (c) Microsoft Corporation.
3 * Licensed under the MIT license.
4 */
5
6//! Multi-vector matrix types and distance functions.
7//!
8//! Row-major matrix abstractions for multi-vector representations, where each
9//! entity is encoded as multiple embedding vectors (e.g., per-token embeddings).
10//!
11//! # Core Types
12//!
13//! | Type | Description |
14//! |------|-------------|
15//! | [`Mat`] | Owning matrix that manages its own memory |
16//! | [`MatRef`] | Immutable borrowed view  |
17//! | [`MatMut`] | Mutable borrowed view |
18//! | [`Repr`] | Trait defining row layout (e.g., [`Standard`]) |
19//! | [`BlockTransposed`] | Owning block-transposed matrix |
20//! | [`BlockTransposedRef`] | Immutable view of a block-transposed matrix |
21//! | [`BlockTransposedMut`] | Mutable view of a block-transposed matrix |
22//! | [`QueryMatRef`] | Query wrapper for asymmetric distances |
23//! | [`QueryComputer`] | Architecture-dispatched SIMD query computer |
24//! | [`MaxSim`] | Per-query-vector max similarity computation |
25//! | [`Chamfer`] | Asymmetric Chamfer distance (sum of MaxSim) |
26//!
27//! # Example
28//!
29//! ```
30//! use diskann_quantization::multi_vector::{
31//!     distance::QueryMatRef,
32//!     Chamfer, Mat, MatMut, MatRef, MaxSim, Standard,
33//! };
34//! use diskann_utils::ReborrowMut;
35//! use diskann_vector::{DistanceFunctionMut, PureDistanceFunction};
36//!
37//! // Create an owned matrix (2 vectors, dim 3, initialized to 0.0)
38//! let mut owned = Mat::new(Standard::new(2, 3).unwrap(), 0.0f32).unwrap();
39//! assert_eq!(owned.num_vectors(), 2);
40//!
41//! // Modify via mutable view
42//! let mut view = owned.reborrow_mut();
43//! if let Some(row) = view.get_row_mut(0) {
44//!    row[0] = 1.0;
45//! }
46//!
47//! // Create views from slices
48//! let query_data = [1.0f32, 0.0, 0.0, 1.0];
49//! let doc_data = [1.0f32, 0.0, 0.0, 1.0];
50//!
51//! // Wrap query as QueryMatRef for type-safe asymmetric distance
52//! let query: QueryMatRef<_> = MatRef::new(
53//!     Standard::new(2, 2).unwrap(),
54//!     &query_data,
55//! ).unwrap().into();
56//! let doc = MatRef::new(Standard::new(2, 2).unwrap(), &doc_data).unwrap();
57//!
58//! // Chamfer distance (sum of max similarities)
59//! let distance = Chamfer::evaluate(query, doc);
60//! assert_eq!(distance, -2.0); // Perfect match: -1.0 per vector
61//!
62//! // MaxSim (per-query-vector scores)
63//! let mut scores = vec![0.0f32; 2];
64//! let mut max_sim = MaxSim::new(&mut scores).unwrap();
65//! max_sim.evaluate(query, doc);
66//! assert_eq!(scores[0], -1.0);
67//! assert_eq!(scores[1], -1.0);
68//! ```
69
70pub mod block_transposed;
71pub mod distance;
72pub(crate) mod matrix;
73
74pub use block_transposed::{BlockTransposed, BlockTransposedMut, BlockTransposedRef};
75pub use distance::{Chamfer, MaxSim, MaxSimError, QueryComputer, QueryMatRef};
76pub use matrix::{
77    Defaulted, LayoutError, Mat, MatMut, MatRef, NewCloned, NewMut, NewOwned, NewRef, Overflow,
78    Repr, ReprMut, ReprOwned, SliceError, Standard,
79};
diskann_quantization/multi_vector/mod.rs

diskann_quantization/multi_vector/
mod.rs
