fix input dtype

Signed-off-by: Connor Tsui <connor.tsui20@gmail.com>

Author: connortsui20

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

@@ -15,12 +15,15 @@ use vortex_array::ArrayView;
 use vortex_array::ExecutionCtx;
 use vortex_array::IntoArray;
 use vortex_array::arrays::Extension;
+use vortex_array::arrays::ExtensionArray;
 use vortex_array::arrays::FixedSizeListArray;
 use vortex_array::arrays::PrimitiveArray;
 use vortex_array::arrays::dict::DictArray;
 use vortex_array::arrays::extension::ExtensionArrayExt;
 use vortex_array::arrays::fixed_size_list::FixedSizeListArrayExt;
 use vortex_array::dtype::Nullability;
+use vortex_array::dtype::extension::ExtDType;
+use vortex_array::extension::EmptyMetadata;
 use vortex_array::validity::Validity;
 use vortex_buffer::BufferMut;
 use vortex_error::VortexExpect;
@@ -38,6 +41,7 @@ use crate::scalar_fns::sorf_transform::SorfOptions;
 use crate::scalar_fns::sorf_transform::SorfTransform;
 use crate::utils::cast_to_f32;
 use crate::vector::AnyVector;
+use crate::vector::Vector;
 
 /// Configuration for TurboQuant encoding.
 #[derive(Clone, Debug)]
@@ -236,6 +240,7 @@ pub unsafe fn turboquant_encode_unchecked(
             Validity::NonNullable,
             0,
         )?;
+        let empty_padded_vector = wrap_padded_as_vector(empty_fsl.into_array())?;
 
         let sorf_options = SorfOptions {
             seed,
@@ -244,19 +249,27 @@ pub unsafe fn turboquant_encode_unchecked(
             element_ptype,
         };
         return Ok(
-            SorfTransform::try_new_array(&sorf_options, empty_fsl.into_array(), 0)?.into_array(),
+            SorfTransform::try_new_array(&sorf_options, empty_padded_vector, 0)?.into_array(),
         );
     }
 
     let core = turboquant_quantize_core(&fsl, seed, config.bit_width, config.num_rounds, ctx)?;
     let quantized_fsl =
         build_quantized_fsl(num_rows, core.all_indices, &core.centroids, core.padded_dim)?;
+    let padded_vector = wrap_padded_as_vector(quantized_fsl)?;
 
     let sorf_options = SorfOptions {
         seed,
         num_rounds: config.num_rounds,
         dimension,
         element_ptype,
     };
-    Ok(SorfTransform::try_new_array(&sorf_options, quantized_fsl, num_rows)?.into_array())
+    Ok(SorfTransform::try_new_array(&sorf_options, padded_vector, num_rows)?.into_array())
+}
+
+/// Wrap an `FSL<f32, padded_dim>` in a [`Vector`](crate::vector::Vector) extension so it can be
+/// passed as the child of [`SorfTransform`], which expects a `Vector<padded_dim>` input.
+fn wrap_padded_as_vector(fsl: ArrayRef) -> VortexResult<ArrayRef> {
+    let ext_dtype = ExtDType::<Vector>::try_new(EmptyMetadata, fsl.dtype().clone())?.erased();
+    Ok(ExtensionArray::new(ext_dtype, fsl).into_array())
 }

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

@@ -128,10 +128,11 @@ fn unwrap_codes_centroids_norms(
     ctx: &mut vortex_array::ExecutionCtx,
 ) -> VortexResult<(PrimitiveArray, PrimitiveArray, PrimitiveArray)> {
     let (sorf_child, norms_child) = unwrap_l2denorm(encoded);
-    let fsl_child = unwrap_sorf(&sorf_child);
+    let padded_vector_child = unwrap_sorf(&sorf_child);
 
-    // FSL(Dict(codes, centroids))
-    let fsl: FixedSizeListArray = fsl_child.execute(ctx)?;
+    // Vector<padded_dim> wrapping FSL(Dict(codes, centroids))
+    let padded_vector: ExtensionArray = padded_vector_child.execute(ctx)?;
+    let fsl: FixedSizeListArray = padded_vector.storage_array().clone().execute(ctx)?;
     let dict = fsl
         .elements()
         .as_opt::<Dict>()

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

@@ -225,6 +225,8 @@ fn dot_product_quantized_accuracy() -> VortexResult<()> {
 fn sorf_transform_roundtrip_isolation() -> VortexResult<()> {
     use vortex_array::IntoArray;
     use vortex_array::arrays::dict::DictArray;
+    use vortex_array::dtype::extension::ExtDType;
+    use vortex_array::extension::EmptyMetadata;
     use vortex_array::validity::Validity;
     use vortex_buffer::BufferMut;
 
@@ -234,6 +236,7 @@ fn sorf_transform_roundtrip_isolation() -> VortexResult<()> {
     use crate::scalar_fns::sorf_transform::SorfMatrix;
     use crate::scalar_fns::sorf_transform::SorfOptions;
     use crate::scalar_fns::sorf_transform::SorfTransform;
+    use crate::vector::Vector;
 
     let dim = 128usize;
     let seed = 99u64;
@@ -287,14 +290,20 @@ fn sorf_transform_roundtrip_isolation() -> VortexResult<()> {
         num_rows,
     )?;
 
+    // Wrap the padded FSL in a Vector extension so it can be the SorfTransform child.
+    let padded_vector_dtype =
+        ExtDType::<Vector>::try_new(EmptyMetadata, fsl.dtype().clone())?.erased();
+    let padded_vector = ExtensionArray::new(padded_vector_dtype, fsl.into_array());
+
     // Wrap in SorfTransform and execute.
     let sorf_options = SorfOptions {
         seed,
         num_rounds,
         dimension: dim as u32,
         element_ptype: vortex_array::dtype::PType::F32,
     };
-    let sorf_array = SorfTransform::try_new_array(&sorf_options, fsl.into_array(), num_rows)?;
+    let sorf_array =
+        SorfTransform::try_new_array(&sorf_options, padded_vector.into_array(), num_rows)?;
 
     let mut ctx = SESSION.create_execution_ctx();
     let result: ExtensionArray = sorf_array.into_array().execute(&mut ctx)?;

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

@@ -8,17 +8,35 @@
 //! Walsh-Hadamard transform to achieve O(d log d) matrix-vector products instead of the O(d^2) cost
 //! of a dense orthogonal matrix.
 //!
-//! This module wraps an FSL child (e.g. `FSL(Dict(codes, centroids))`) and applies the inverse SORF
-//! transform at execution time, producing a [`Vector`] extension array with the original
-//! (pre-padding) dimensionality.
+//! This module wraps a [`Vector`] extension array whose dimension is the padded SORF dimension
+//! (e.g. a `Vector` wrapping `FSL(Dict(codes, centroids))`) and applies the inverse SORF transform
+//! at execution time, producing a [`Vector`] extension array with the original (pre-padding)
+//! dimensionality.
 //!
 //! The transform parameters are stored as a deterministic seed in [`SorfOptions`], so the
 //! [`SorfMatrix`] is reconstructed cheaply at decode time. Sign diagonals are defined by Vortex's
 //! frozen local SplitMix64 stream contract rather than by an external RNG crate.
 //!
-//! **All SORF computation happens in f32.** Input elements of other float types (f16, f64) are cast
-//! to f32 before the transform, and the result is cast back to the target type specified by
-//! [`SorfOptions::element_ptype`].
+//! # Input element type: `f32` only (TODO(connor): for now...)
+//!
+//! The child [`Vector`] **must** have `f32` storage elements. This is a hard constraint that is
+//! enforced by `SorfTransform`'s `return_dtype` check. Callers with `f16` or `f64` source data need
+//! to cast to `f32` before wrapping in a [`Vector`] and handing it to SorfTransform.
+//!
+//! The reason for this constraint is that TurboQuant (the only production caller today) stores its
+//! dictionary centroids as `f32`, and the SORF transform itself operates internally in `f32`.
+//!
+//! Supporting other float storage types would require an implicit up-/down-cast that we do not yet
+//! want to bake into SorfTransform. This restriction is intentional and may be relaxed in the
+//! future, but today it is load-bearing.
+//!
+//! # Output element type
+//!
+//! The output [`Vector`]'s element type is whatever [`SorfOptions::element_ptype`] is set to. It
+//! does **not** have to match the child's `f32` storage: we apply an explicit `f32 -> T` cast
+//! while materializing the output. This lets SorfTransform hand its result directly to a
+//! downstream consumer (e.g. [`L2Denorm`](crate::scalar_fns::l2_denorm::L2Denorm)) whose
+//! element-type expectation may differ from the `f32` the transform operated on internally.
 //!
 //! [sorf-paper]: https://proceedings.neurips.cc/paper_files/paper/2016/file/53adaf494dc89ef7196d73636eb2451b-Paper.pdf
 //! [`Vector`]: crate::vector::Vector
@@ -41,9 +59,13 @@ mod vtable;
 
 /// Inverse SORF orthogonal transform scalar function.
 ///
-/// Applies the inverse structured Walsh-Hadamard orthogonal transform to an FSL child,
-/// truncates from padded dimension to the original dimension, casts to the target element
-/// type, and wraps in a [`Vector`](crate::vector::Vector) extension array.
+/// Takes a [`Vector`](crate::vector::Vector) extension child at the padded dimension with `f32`
+/// storage, applies the inverse structured Walsh-Hadamard orthogonal transform, truncates to the
+/// original (pre-padding) dimension, casts element-wise to [`SorfOptions::element_ptype`], and
+/// wraps the result in a new [`Vector`](crate::vector::Vector) extension array.
+///
+/// See the [module-level docs](crate::scalar_fns::sorf_transform) for the rationale behind the
+/// `f32`-only input constraint.
 #[derive(Clone)]
 pub struct SorfTransform;
 
@@ -57,9 +79,12 @@ pub struct SorfOptions {
     pub seed: u64,
     /// Number of sign-diagonal + WHT rounds in the structured orthogonal transform.
     pub num_rounds: u8,
-    /// Original vector dimension (before power-of-2 padding).
+    /// Original vector dimension (before power-of-2 padding). The output
+    /// [`Vector`](crate::vector::Vector) has this dimension.
     pub dimension: u32,
-    /// Target output element type (e.g. `F16`, `F32`, `F64`).
+    /// Element type of the output [`Vector`](crate::vector::Vector). The child input must always
+    /// be `f32`, but the output can be any float type (`F16`, `F32`, `F64`); the final
+    /// `f32 -> element_ptype` cast happens while building the output.
     pub element_ptype: PType,
 }
 
@@ -71,8 +96,16 @@ impl SorfTransform {
 
     /// Constructs a validated [`ScalarFnArray`] that lazily applies the inverse SORF transform.
     ///
-    /// The `child` must be a `FixedSizeList` (or array that executes to one) with logical float
-    /// elements and `list_size == padded_dim` (i.e. `dimension.next_power_of_two()`).
+    /// The `child` must be a [`Vector`] extension array (or an array that executes to one) with:
+    ///
+    /// - dimension equal to `padded_dim` (i.e. `options.dimension.next_power_of_two()`), and
+    /// - `f32` storage elements. This is a hard requirement today; see the
+    ///   [module-level docs](crate::scalar_fns::sorf_transform) for the rationale.
+    ///
+    /// The output [`Vector`] has dimension `options.dimension` and element type
+    /// `options.element_ptype`.
+    ///
+    /// [`Vector`]: crate::vector::Vector
     pub fn try_new_array(
         options: &SorfOptions,
         child: ArrayRef,