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//! # Example
12//!
13//! ```
14//! use diskann_quantization::multi_vector::{
15//! distance::QueryMatRef,
16//! Chamfer, Mat, MatMut, MatRef, MaxSim, Standard,
17//! };
18//! use diskann_utils::ReborrowMut;
19//! use diskann_vector::{DistanceFunctionMut, PureDistanceFunction};
20//!
21//! // Create an owned matrix (2 vectors, dim 3, initialized to 0.0)
22//! let mut owned = Mat::new(Standard::new(2, 3).unwrap(), 0.0f32).unwrap();
23//! assert_eq!(owned.num_vectors(), 2);
24//!
25//! // Modify via mutable view
26//! let mut view = owned.reborrow_mut();
27//! if let Some(row) = view.get_row_mut(0) {
28//! row[0] = 1.0;
29//! }
30//!
31//! // Create views from slices
32//! let query_data = [1.0f32, 0.0, 0.0, 1.0];
33//! let doc_data = [1.0f32, 0.0, 0.0, 1.0];
34//!
35//! // Wrap query as QueryMatRef for type-safe asymmetric distance
36//! let query: QueryMatRef<_> = MatRef::new(
37//! Standard::new(2, 2).unwrap(),
38//! &query_data,
39//! ).unwrap().into();
40//! let doc = MatRef::new(Standard::new(2, 2).unwrap(), &doc_data).unwrap();
41//!
42//! // Chamfer distance (sum of max similarities)
43//! let distance = Chamfer::evaluate(query, doc);
44//! assert_eq!(distance, -2.0); // Perfect match: -1.0 per vector
45//!
46//! // MaxSim (per-query-vector scores)
47//! let mut scores = vec![0.0f32; 2];
48//! let mut max_sim = MaxSim::new(&mut scores);
49//! max_sim.evaluate(query, doc);
50//! assert_eq!(scores[0], -1.0);
51//! assert_eq!(scores[1], -1.0);
52//! ```
53
54pub mod block_transposed;
55pub mod distance;
56pub(crate) mod matrix;
57
58pub use block_transposed::{BlockTransposed, BlockTransposedMut, BlockTransposedRef};
59pub use distance::{
60 BoxErase, Chamfer, Erase, MaxSim, MaxSimElement, MaxSimError, MaxSimIsa, MaxSimKernel,
61 NotSupported, ProjectedEigen, QueryMatRef, build_max_sim,
62};
63pub use matrix::{
64 Defaulted, LayoutError, Mat, MatMut, MatRef, NewCloned, NewMut, NewOwned, NewRef, Overflow,
65 Repr, ReprMut, ReprOwned, SliceError, Standard,
66};