nabled_embeddings/
normalize.rs

1//! Row-wise L2 normalization for embedding matrices.
2//!
3//! Normalizing rows to unit length is the canonical pre-step for cosine retrieval: once both
4//! query and corpus are normalized, a plain dot product equals cosine similarity. These functions
5//! delegate to `nabled_linalg::vector::batched_normalize*` and add nothing but the embedding-domain
6//! error type. Zero-norm rows are left at the origin (divided by a small epsilon) rather than
7//! erroring, matching the underlying kernel.
8
9use nabled_core::scalar::NabledReal;
10use nabled_linalg::vector;
11use ndarray::{Array2, ArrayBase, ArrayView2, Data, DataMut, Ix2};
12
13use crate::error::EmbeddingError;
14
15/// Normalize each row of `rows` to unit L2 length.
16///
17/// # Errors
18/// Returns [`EmbeddingError::EmptyInput`] when `rows` is empty.
19pub fn normalize_rows<T: NabledReal>(rows: &Array2<T>) -> Result<Array2<T>, EmbeddingError> {
20    Ok(vector::batched_normalize(rows)?)
21}
22
23/// Normalize each row of a matrix view to unit L2 length.
24///
25/// # Errors
26/// Returns [`EmbeddingError::EmptyInput`] when `rows` is empty.
27pub fn normalize_rows_view<T: NabledReal>(
28    rows: &ArrayView2<'_, T>,
29) -> Result<Array2<T>, EmbeddingError> {
30    Ok(vector::batched_normalize_view(rows)?)
31}
32
33/// Normalize each row into a caller-provided `output` buffer.
34///
35/// # Errors
36/// Returns [`EmbeddingError::EmptyInput`] when `rows` is empty or
37/// [`EmbeddingError::DimensionMismatch`] when `output` is not the same shape as `rows`.
38pub fn normalize_rows_into<T, S>(
39    rows: &ArrayBase<S, Ix2>,
40    output: &mut ArrayBase<impl DataMut<Elem = T>, Ix2>,
41) -> Result<(), EmbeddingError>
42where
43    T: NabledReal,
44    S: Data<Elem = T>,
45{
46    vector::batched_normalize_into(rows, output)?;
47    Ok(())
48}
49
50#[cfg(test)]
51mod tests {
52    use ndarray::{Array2, arr2};
53
54    use super::*;
55
56    fn unit_norm<T: NabledReal>(row: &ndarray::ArrayView1<'_, T>) -> T {
57        row.iter().copied().fold(T::zero(), |acc, v| acc + v * v).sqrt()
58    }
59
60    #[test]
61    fn normalize_rows_yields_unit_rows_f64() {
62        let rows = arr2(&[[3.0_f64, 4.0], [1.0, 0.0]]);
63        let normalized = normalize_rows(&rows).unwrap();
64        assert!((unit_norm(&normalized.row(0)) - 1.0).abs() < 1e-12);
65        assert!((normalized[[0, 0]] - 0.6).abs() < 1e-12);
66        assert!((normalized[[0, 1]] - 0.8).abs() < 1e-12);
67    }
68
69    #[test]
70    fn normalize_rows_yields_unit_rows_f32() {
71        let rows = arr2(&[[3.0_f32, 4.0]]);
72        let normalized = normalize_rows(&rows).unwrap();
73        assert!((unit_norm(&normalized.row(0)) - 1.0).abs() < 1e-5);
74    }
75
76    #[test]
77    fn normalize_rows_view_matches_owned() {
78        let rows = arr2(&[[3.0_f64, 4.0], [0.0, 2.0]]);
79        let owned = normalize_rows(&rows).unwrap();
80        let viewed = normalize_rows_view(&rows.view()).unwrap();
81        assert_eq!(owned, viewed);
82    }
83
84    #[test]
85    fn normalize_rows_into_matches_allocating() {
86        let rows = arr2(&[[3.0_f64, 4.0], [0.0, 2.0]]);
87        let expected = normalize_rows(&rows).unwrap();
88        let mut output = Array2::<f64>::zeros(rows.dim());
89        normalize_rows_into(&rows, &mut output).unwrap();
90        assert_eq!(output, expected);
91    }
92
93    #[test]
94    fn normalize_rows_into_rejects_shape_mismatch() {
95        let rows = arr2(&[[3.0_f64, 4.0]]);
96        let mut output = Array2::<f64>::zeros((2, 2));
97        assert_eq!(normalize_rows_into(&rows, &mut output), Err(EmbeddingError::DimensionMismatch));
98    }
99
100    #[test]
101    fn normalize_rows_rejects_empty() {
102        let rows = Array2::<f64>::zeros((0, 0));
103        assert_eq!(normalize_rows(&rows), Err(EmbeddingError::EmptyInput));
104    }
105
106    #[test]
107    fn normalize_rows_keeps_zero_row_finite() {
108        let rows = arr2(&[[0.0_f64, 0.0]]);
109        let normalized = normalize_rows(&rows).unwrap();
110        assert!(normalized.iter().all(|v| v.is_finite()));
111    }
112}
nabled_embeddings/normalize.rs

nabled_embeddings/
normalize.rs
