| 208 | } |
| 209 | |
| 210 | fn pq_encode(v: &[f32], codebooks: &[Vec<Vec<f32>>], m: usize, sub_dim: usize) -> Vec<u8> { |
| 211 | let mut codes = Vec::with_capacity(m); |
| 212 | #[allow(clippy::needless_range_loop)] |
| 213 | for s in 0..m { |
| 214 | let offset = s * sub_dim; |
| 215 | let sub = &v[offset..offset + sub_dim]; |
| 216 | let best = codebooks[s] |
| 217 | .iter() |
| 218 | .enumerate() |
| 219 | .min_by(|(_, a), (_, b)| { |
| 220 | l2_sq(sub, a) |
| 221 | .partial_cmp(&l2_sq(sub, b)) |
| 222 | .unwrap_or(std::cmp::Ordering::Equal) |
| 223 | }) |
| 224 | .map(|(i, _)| i) |
| 225 | .unwrap_or(0); |
| 226 | codes.push(best as u8); |
| 227 | } |
| 228 | codes |
| 229 | } |
| 230 | |
| 231 | fn train_codebooks( |
| 232 | rotated: &[Vec<f32>], |