Small cleanups.

Author: Mark Hildebrand

diskann-quantization/src/product/tables/basic.rs

@@ -119,23 +119,23 @@ where
 
     /// Return a [`Strided`] for the raw pivots of the requested chunk.
     ///
-    /// Returns `None` if `chunk >= self.nchunks`.
+    /// Returns `None` if `chunk >= self.nchunks()`.
     pub fn pivots_for(&self, chunk: usize) -> Option<Strided<'_, f32>> {
         let range = self.offsets.get(chunk)?;
 
         #[expect(
             clippy::expect_used,
             reason = "the BasicTable's invariants mean this panic should be unreachable"
         )]
-        Some(
-            Strided::try_from_data(
-                &self.pivots.as_slice()[range.start..],
-                self.pivots.nrows(),
-                range.len(),
-                self.pivots.ncols(),
-            )
-            .expect("BasicTable asserts that this layout is valid"),
+        let strided = Strided::try_from_data(
+            &self.pivots.as_slice()[range.start..],
+            self.pivots.nrows(),
+            range.len(),
+            self.pivots.ncols(),
         )
+        .expect("BasicTable asserts that this layout is valid");
+
+        Some(strided)
     }
 }
 
@@ -329,7 +329,7 @@ mod tests {
                 let (pivots, offsets) = create_pivot_tables(schema.to_owned(), num_centers);
                 let table = BasicTable::new(pivots, offsets).unwrap();
 
-                // Check that `pivots_for` works as expected with repsect to the documented
+                // Check that `pivots_for` works as expected with respect to the documented
                 // table configuration for `create_pivot_tables`.
                 for chunk in 0..schema.len() {
                     let strided = table.pivots_for(chunk).unwrap();
@@ -340,20 +340,17 @@ mod tests {
                         let base = ((center + chunk) % num_centers) as f32;
                         row.iter().enumerate().for_each(|(dim, b)| {
                             let offset = if dim.is_multiple_of(2) { 0.25 } else { -0.25 };
-
-                            if dim.is_multiple_of(2) {
-                                assert_eq!(
-                                    *b,
-                                    base + offset,
-                                    "failed: chunk {} of {}, center {} of {}, dim {} of {}",
-                                    chunk,
-                                    schema.len(),
-                                    center,
-                                    num_centers,
-                                    dim,
-                                    row.len(),
-                                );
-                            }
+                            assert_eq!(
+                                *b,
+                                base + offset,
+                                "failed: chunk {} of {}, center {} of {}, dim {} of {}",
+                                chunk,
+                                schema.len(),
+                                center,
+                                num_centers,
+                                dim,
+                                row.len(),
+                            );
                         })
                     }
                 }

diskann-quantization/src/product/train.rs

@@ -62,7 +62,7 @@ impl TrainQuantizer for LightPQTrainingParameters {
     type Error = PQTrainingError;
 
     /// Perform product quantization training on the provided training set and return a
-    /// `SimplePivots` containing the result of kmeans clustering on each partition.
+    /// [`BasicTable`] containing the result of kmeans clustering on each partition.
     ///
     /// # Panics
     ///
@@ -208,7 +208,7 @@ impl TrainQuantizer for LightPQTrainingParameters {
                 BasicTable::new(packed, schema.to_owned()).map_err(|err| PQTrainingError {
                     chunk: schema.len(),
                     of: schema.len(),
-                    dim: data.nrows(),
+                    dim: data.ncols(),
                     kind: PQTrainingErrorKind::InternalError(Box::new(err)),
                 })?;