TurboQuant again!

Cargo.toml

@@ -13,6 +13,7 @@ members = [
     "vortex-proto",
     "vortex-array",
     "vortex-tensor",
+    "vortex-turboquant",
     "vortex-compressor",
     "vortex-btrblocks",
     "vortex-layout",
@@ -296,6 +297,7 @@ vortex-sequence = { version = "0.1.0", path = "encodings/sequence", default-feat
 vortex-session = { version = "0.1.0", path = "./vortex-session", default-features = false }
 vortex-sparse = { version = "0.1.0", path = "./encodings/sparse", default-features = false }
 vortex-tensor = { version = "0.1.0", path = "./vortex-tensor", default-features = false }
+vortex-turboquant = { version = "0.1.0", path = "./vortex-turboquant", default-features = false }
 vortex-utils = { version = "0.1.0", path = "./vortex-utils", default-features = false }
 vortex-zigzag = { version = "0.1.0", path = "./encodings/zigzag", default-features = false }
 vortex-zstd = { version = "0.1.0", path = "./encodings/zstd", default-features = false }

vortex-tensor/public-api.lock

@@ -528,6 +528,10 @@ pub fn vortex_tensor::vector::AnyVector::try_match<'a>(&'a vortex_array::dtype::
 
 pub struct vortex_tensor::vector::Vector
 
+impl vortex_tensor::vector::Vector
+
+pub fn vortex_tensor::vector::Vector::try_new_vector_array(vortex_array::array::erased::ArrayRef) -> vortex_error::VortexResult<vortex_array::array::erased::ArrayRef>
+
 impl core::clone::Clone for vortex_tensor::vector::Vector
 
 pub fn vortex_tensor::vector::Vector::clone(&self) -> vortex_tensor::vector::Vector

vortex-tensor/src/encodings/turboquant/compress.rs

@@ -205,8 +205,8 @@ fn turboquant_quantize_core(
     let dimension = fsl.list_size() as usize;
     let num_rows = fsl.len();
 
-    let rotation = SorfMatrix::try_new(seed, dimension, num_rounds as usize)?;
-    let padded_dim = rotation.padded_dim();
+    let padded_dim = dimension.next_power_of_two();
+    let rotation = SorfMatrix::try_new_padded(padded_dim, num_rounds as usize, seed)?;
     let padded_dim_u32 =
         u32::try_from(padded_dim).vortex_expect("padded_dim stays representable as u32");
 

vortex-tensor/src/encodings/turboquant/tests/structural.rs

@@ -259,8 +259,8 @@ fn sorf_transform_roundtrip_isolation() -> VortexResult<()> {
     }
 
     // Forward transform + quantize (mimicking what turboquant_quantize_core does).
-    let rotation = SorfMatrix::try_new(seed, dim, num_rounds as usize)?;
-    let padded_dim = rotation.padded_dim();
+    let padded_dim = dim.next_power_of_two();
+    let rotation = SorfMatrix::try_new_padded(padded_dim, num_rounds as usize, seed)?;
     let centroids = compute_or_get_centroids(padded_dim as u32, 8)?;
     let boundaries = compute_centroid_boundaries(&centroids);
 

vortex-tensor/src/scalar_fns/inner_product.rs

@@ -366,7 +366,7 @@ impl InnerProduct {
         let mut padded_query = vec![0.0f32; padded_dim];
         padded_query[..dim].copy_from_slice(flat.as_slice::<f32>());
 
-        let rotation = SorfMatrix::try_new(seed, dim, num_rounds)?;
+        let rotation = SorfMatrix::try_new_padded(padded_dim, num_rounds, seed)?;
         let mut rotated_query = vec![0.0f32; padded_dim];
         rotation.rotate(&padded_query, &mut rotated_query);
 
@@ -930,7 +930,7 @@ mod tests {
             seed: u64,
             num_rounds: u8,
         ) -> VortexResult<Vec<f32>> {
-            let rotation = SorfMatrix::try_new(seed, dim, num_rounds as usize)?;
+            let rotation = SorfMatrix::try_new_padded(padded_dim, num_rounds as usize, seed)?;
             let mut padded = vec![0.0f32; padded_dim];
             let mut rotated = vec![0.0f32; padded_dim];
             let mut out = Vec::with_capacity(num_rows * dim);

vortex-tensor/src/scalar_fns/sorf_transform/rotation.rs

@@ -77,9 +77,25 @@ impl SorfMatrix {
     /// round-major, block-major order, with each `u64` contributing 64 sign bits in
     /// least-significant-bit-first order.
     pub fn try_new(seed: u64, dimensions: usize, num_rounds: usize) -> VortexResult<Self> {
+        Self::try_new_padded(dimensions.next_power_of_two(), num_rounds, seed)
+    }
+
+    /// Create a new structured Walsh-Hadamard-based orthogonal transform for a padded dimension.
+    ///
+    /// `padded_dimensions` must already be a power of two. Callers that start from an unpadded
+    /// logical dimension should call [`Self::try_new`] instead.
+    pub(crate) fn try_new_padded(
+        padded_dimensions: usize,
+        num_rounds: usize,
+        seed: u64,
+    ) -> VortexResult<Self> {
         vortex_ensure!(num_rounds >= 1, "num_rounds must be >= 1, got {num_rounds}");
+        vortex_ensure!(
+            padded_dimensions.is_power_of_two(),
+            "padded_dimensions must be a power of two, got {padded_dimensions}"
+        );
 
-        let padded_dim = dimensions.next_power_of_two();
+        let padded_dim = padded_dimensions;
         let sign_masks = gen_sign_masks_from_seed(seed, padded_dim, num_rounds);
 
         // Compute in f64 for precision, then store as f32 since the WHT operates on f32 buffers.
@@ -132,8 +148,7 @@ impl SorfMatrix {
     /// Apply the forward structured transform: `norm · H · D_k · ... · H · D₁ · x`.
     fn apply_srht(&self, buf: &mut [f32]) {
         for round in 0..self.num_rounds {
-            let offset = round * self.padded_dim;
-            apply_signs_xor(buf, &self.sign_masks[offset..offset + self.padded_dim]);
+            self.apply_signs_xor(buf, round);
             walsh_hadamard_transform(buf);
         }
 
@@ -148,14 +163,24 @@ impl SorfMatrix {
     fn apply_inverse_srht(&self, buf: &mut [f32]) {
         for round in (0..self.num_rounds).rev() {
             walsh_hadamard_transform(buf);
-            let offset = round * self.padded_dim;
-            apply_signs_xor(buf, &self.sign_masks[offset..offset + self.padded_dim]);
+            self.apply_signs_xor(buf, round);
         }
 
         let norm = self.norm_factor;
         buf.iter_mut().for_each(|val| *val *= norm);
     }
 
+    /// Apply one round's sign masks via XOR on the IEEE 754 sign bit.
+    ///
+    /// This is branchless and auto-vectorizes into `vpxor` (x86) / `veor` (ARM). Equivalent to
+    /// multiplying each element by +/-1.0, but avoids FP dependency chains.
+    fn apply_signs_xor(&self, buf: &mut [f32], round: usize) {
+        let masks = &self.sign_masks[round * self.padded_dim..][..self.padded_dim];
+        for (val, &mask) in buf.iter_mut().zip(masks.iter()) {
+            *val = f32::from_bits(val.to_bits() ^ mask);
+        }
+    }
+
     /// Export the sign vectors as a flat `Vec<u8>` of 0/1 values in inverse application order
     /// `[D_k | ... | D₁]`.
     ///
@@ -263,16 +288,6 @@ fn sign_mask_from_word(word: u64, bit_idx: usize) -> u32 {
     }
 }
 
-/// Apply sign masks via XOR on the IEEE 754 sign bit.
-///
-/// This is branchless and auto-vectorizes into `vpxor` (x86) / `veor` (ARM). Equivalent to
-/// multiplying each element by +/-1.0, but avoids FP dependency chains.
-fn apply_signs_xor(buf: &mut [f32], masks: &[u32]) {
-    for (val, &mask) in buf.iter_mut().zip(masks.iter()) {
-        *val = f32::from_bits(val.to_bits() ^ mask);
-    }
-}
-
 /// In-place Fast Walsh-Hadamard Transform (FWHT), unnormalized and iterative.
 ///
 /// Input length must be a power of 2. Runs in O(n log n) via `log2(n)` stages of `n / 2`
@@ -327,14 +342,24 @@ mod tests {
             .collect()
     }
 
+    fn dim_to_usize(dim: u32) -> usize {
+        usize::try_from(dim).unwrap()
+    }
+
+    fn rounds_to_usize(num_rounds: u8) -> usize {
+        usize::from(num_rounds)
+    }
+
     #[test]
     fn deterministic_from_seed() -> VortexResult<()> {
-        let r1 = SorfMatrix::try_new(42, 64, 3)?;
-        let r2 = SorfMatrix::try_new(42, 64, 3)?;
+        let dim = dim_to_usize(64u32);
+        let num_rounds = rounds_to_usize(3u8);
+        let r1 = SorfMatrix::try_new(42u64, dim, num_rounds)?;
+        let r2 = SorfMatrix::try_new(42u64, dim, num_rounds)?;
         let pd = r1.padded_dim();
 
         let mut input = vec![0.0f32; pd];
-        for i in 0..64 {
+        for i in 0..dim {
             input[i] = i as f32;
         }
         let mut out1 = vec![0.0f32; pd];
@@ -349,41 +374,58 @@ mod tests {
 
     #[test]
     fn export_inverse_signs_matches_golden_words() -> VortexResult<()> {
-        let rot = SorfMatrix::try_new(42, 64, 2)?;
+        let dim = dim_to_usize(64u32);
+        let num_rounds = rounds_to_usize(2u8);
+        let seed = 42u64;
+        let rot = SorfMatrix::try_new(seed, dim, num_rounds)?;
+        let padded_dim = rot.padded_dim();
         let actual = rot.export_inverse_signs_u8();
-        let mut rng = SplitMix64::new(42);
+        let mut rng = SplitMix64::new(seed);
         let round0_word = rng.next_u64();
         let round1_word = rng.next_u64();
 
-        let mut expected = Vec::with_capacity(128);
-        expected.extend(unpack_sign_bits(round1_word, 64));
-        expected.extend(unpack_sign_bits(round0_word, 64));
+        let mut expected = Vec::with_capacity(num_rounds * padded_dim);
+        expected.extend(unpack_sign_bits(round1_word, padded_dim));
+        expected.extend(unpack_sign_bits(round0_word, padded_dim));
 
         assert_eq!(actual, expected);
         Ok(())
     }
 
     #[test]
     fn one_word_generates_64_signs_lsb_first() {
-        let masks = gen_sign_masks_from_seed(42, 64, 1);
-        assert_eq!(masks.len(), 64);
+        let seed = 42u64;
+        let padded_dim = dim_to_usize(64u32);
+        let num_rounds = rounds_to_usize(1u8);
+        let masks = gen_sign_masks_from_seed(seed, padded_dim, num_rounds);
+        assert_eq!(masks.len(), padded_dim);
 
-        let mut rng = SplitMix64::new(42);
+        let mut rng = SplitMix64::new(seed);
         let word = rng.next_u64();
-        let expected: Vec<_> = (0..64)
+        let expected: Vec<_> = (0..padded_dim)
             .map(|bit_idx| sign_mask_from_word(word, bit_idx))
             .collect();
         assert_eq!(masks, expected);
     }
 
+    #[test]
+    fn accepts_non_power_of_two_dimensions() -> VortexResult<()> {
+        let rot = SorfMatrix::try_new(42u64, dim_to_usize(100u32), rounds_to_usize(3u8))?;
+        assert_eq!(rot.padded_dim(), 128);
+        Ok(())
+    }
+
     #[test]
     fn tail_block_uses_only_required_bits() {
-        let masks = gen_sign_masks_from_seed(42, 32, 1);
-        assert_eq!(masks.len(), 32);
+        let seed = 42u64;
+        let padded_dim = dim_to_usize(32u32);
+        let num_rounds = rounds_to_usize(1u8);
+        let masks = gen_sign_masks_from_seed(seed, padded_dim, num_rounds);
+        assert_eq!(masks.len(), padded_dim);
 
-        let mut rng = SplitMix64::new(42);
+        let mut rng = SplitMix64::new(seed);
         let word = rng.next_u64();
-        let expected: Vec<_> = (0..32)
+        let expected: Vec<_> = (0..padded_dim)
             .map(|bit_idx| sign_mask_from_word(word, bit_idx))
             .collect();
         assert_eq!(masks, expected);
@@ -392,19 +434,21 @@ mod tests {
     /// Verify roundtrip is exact to f32 precision across many dimensions and round counts,
     /// including non-power-of-two dimensions that require padding.
     #[rstest]
-    #[case(32, 3)]
-    #[case(64, 3)]
-    #[case(100, 3)]
-    #[case(128, 1)]
-    #[case(128, 2)]
-    #[case(128, 3)]
-    #[case(128, 5)]
-    #[case(256, 3)]
-    #[case(512, 3)]
-    #[case(768, 3)]
-    #[case(1024, 3)]
-    fn roundtrip_exact(#[case] dim: usize, #[case] num_rounds: usize) -> VortexResult<()> {
-        let rot = SorfMatrix::try_new(42, dim, num_rounds)?;
+    #[case(32u32, 3u8)]
+    #[case(64u32, 3u8)]
+    #[case(100u32, 3u8)]
+    #[case(128u32, 1u8)]
+    #[case(128u32, 2u8)]
+    #[case(128u32, 3u8)]
+    #[case(128u32, 5u8)]
+    #[case(256u32, 3u8)]
+    #[case(512u32, 3u8)]
+    #[case(768u32, 3u8)]
+    #[case(1024u32, 3u8)]
+    fn roundtrip_exact(#[case] dim: u32, #[case] num_rounds: u8) -> VortexResult<()> {
+        let dim = dim_to_usize(dim);
+        let num_rounds = rounds_to_usize(num_rounds);
+        let rot = SorfMatrix::try_new(42u64, dim, num_rounds)?;
         let padded_dim = rot.padded_dim();
 
         let mut input = vec![0.0f32; padded_dim];
@@ -435,12 +479,14 @@ mod tests {
 
     /// Verify norm preservation across dimensions and round counts.
     #[rstest]
-    #[case(128, 1)]
-    #[case(128, 3)]
-    #[case(128, 5)]
-    #[case(768, 3)]
-    fn preserves_norm(#[case] dim: usize, #[case] num_rounds: usize) -> VortexResult<()> {
-        let rot = SorfMatrix::try_new(7, dim, num_rounds)?;
+    #[case(128u32, 1u8)]
+    #[case(128u32, 3u8)]
+    #[case(128u32, 5u8)]
+    #[case(768u32, 3u8)]
+    fn preserves_norm(#[case] dim: u32, #[case] num_rounds: u8) -> VortexResult<()> {
+        let dim = dim_to_usize(dim);
+        let num_rounds = rounds_to_usize(num_rounds);
+        let rot = SorfMatrix::try_new(42u64, dim, num_rounds)?;
         let padded_dim = rot.padded_dim();
 
         let mut input = vec![0.0f32; padded_dim];
@@ -465,16 +511,15 @@ mod tests {
 
     /// Verify that export -> [`from_u8_slice`] produces identical transform output.
     #[rstest]
-    #[case(64, 3)]
-    #[case(128, 1)]
-    #[case(128, 3)]
-    #[case(128, 5)]
-    #[case(768, 3)]
-    fn sign_export_import_roundtrip(
-        #[case] dim: usize,
-        #[case] num_rounds: usize,
-    ) -> VortexResult<()> {
-        let rot = SorfMatrix::try_new(42, dim, num_rounds)?;
+    #[case(64u32, 3u8)]
+    #[case(128u32, 1u8)]
+    #[case(128u32, 3u8)]
+    #[case(128u32, 5u8)]
+    #[case(768u32, 3u8)]
+    fn sign_export_import_roundtrip(#[case] dim: u32, #[case] num_rounds: u8) -> VortexResult<()> {
+        let dim = dim_to_usize(dim);
+        let num_rounds = rounds_to_usize(num_rounds);
+        let rot = SorfMatrix::try_new(42u64, dim, num_rounds)?;
         let padded_dim = rot.padded_dim();
 
         let signs_u8 = rot.export_inverse_signs_u8();

vortex-tensor/src/scalar_fns/sorf_transform/tests.rs

@@ -65,8 +65,8 @@ fn forward_rotate_and_quantize(
         }
     }
 
-    let rotation = SorfMatrix::try_new(seed, dim, num_rounds)?;
-    let padded_dim = rotation.padded_dim();
+    let padded_dim = dim.next_power_of_two();
+    let rotation = SorfMatrix::try_new_padded(padded_dim, num_rounds, seed)?;
     let centroids = compute_or_get_centroids(padded_dim as u32, bit_width)?;
     let boundaries = compute_centroid_boundaries(&centroids);
 

vortex-tensor/src/scalar_fns/sorf_transform/vtable.rs

@@ -164,7 +164,8 @@ impl ScalarFnVTable for SorfTransform {
         let f32_elements = elements_prim.into_buffer::<f32>();
 
         // Reconstruct the orthogonal transform matrix from the seed.
-        let rotation = SorfMatrix::try_new(options.seed, dim, options.num_rounds as usize)?;
+        let rotation =
+            SorfMatrix::try_new_padded(padded_dim, options.num_rounds as usize, options.seed)?;
 
         // Inverse transform each row, truncate to original dimension, cast to target type.
         match_each_float_ptype!(options.element_ptype, |T| {

vortex-tensor/src/types/vector/mod.rs

@@ -49,7 +49,7 @@ impl Vector {
     /// # Errors
     ///
     /// Returns an error if the [`Vector`] extension dtype rejects the storage array.
-    pub(crate) fn try_new_vector_array(storage: ArrayRef) -> VortexResult<ArrayRef> {
+    pub fn try_new_vector_array(storage: ArrayRef) -> VortexResult<ArrayRef> {
         ExtensionArray::try_new_from_vtable(Vector, EmptyMetadata, storage)
             .map(|ext| ext.into_array())
     }