Make normalized tensor shortcuts explicit

vortex-tensor/benches/cosine_similarity.rs

@@ -28,6 +28,8 @@ use vortex_array::dtype::PType;
 use vortex_array::scalar::Scalar;
 use vortex_array::validity::Validity;
 use vortex_buffer::Buffer;
+use vortex_tensor::encodings::normalized::Normalized;
+use vortex_tensor::scalar_fns::NormMode;
 use vortex_tensor::scalar_fns::cosine_similarity::CosineSimilarity;
 use vortex_tensor::vector::Vector;
 
@@ -80,12 +82,37 @@ fn constant_vector(width: usize) -> ArrayRef {
     ConstantArray::new(ext_scalar, ELEMENTS / width).into_array()
 }
 
-fn bench_cosine(bencher: Bencher, lhs: ArrayRef, rhs: ArrayRef) {
+fn normalized_vectors(width: usize, offset: usize) -> ArrayRef {
+    let row_count = ELEMENTS / width;
+    let elements: Buffer<f64> = (0..ELEMENTS)
+        .map(|i| {
+            if i % width == offset % width {
+                1.0
+            } else {
+                0.0
+            }
+        })
+        .collect();
+    let storage = FixedSizeListArray::new(
+        elements.into_array(),
+        u32::try_from(width).unwrap(),
+        Validity::NonNullable,
+        row_count,
+    )
+    .into_array();
+    let direction = Vector::try_new_vector_array(storage).unwrap();
+    let norms = PrimitiveArray::from_iter((1..=row_count).map(|norm| norm as f64)).into_array();
+
+    // SAFETY: Every direction row is unit length, and the non-negative norms have matching length.
+    unsafe { Normalized::new_unchecked(direction, norms, Validity::NonNullable) }.into_array()
+}
+
+fn bench_cosine(bencher: Bencher, lhs: ArrayRef, rhs: ArrayRef, mode: NormMode) {
     let session = vortex_array::array_session();
     bencher
         .with_inputs(|| {
             (
-                CosineSimilarity::try_new(lhs.clone(), rhs.clone())
+                CosineSimilarity::try_new(lhs.clone(), rhs.clone(), mode)
                     .unwrap()
                     .into_array(),
                 session.create_execution_ctx(),
@@ -97,13 +124,23 @@ fn bench_cosine(bencher: Bencher, lhs: ArrayRef, rhs: ArrayRef) {
 /// The control: both operands vary by row, so every norm must be computed in the row loop.
 #[divan::bench(args = WIDTHS)]
 fn column_x_column(bencher: Bencher, width: usize) {
-    bench_cosine(bencher, vectors(width, 0), vectors(width, 31));
+    bench_cosine(
+        bencher,
+        vectors(width, 0),
+        vectors(width, 31),
+        NormMode::Exact,
+    );
 }
 
 /// The rhs is a broadcast query vector, whose norm is the same in every row.
 #[divan::bench(args = WIDTHS)]
 fn column_x_constant(bencher: Bencher, width: usize) {
-    bench_cosine(bencher, vectors(width, 0), constant_vector(width));
+    bench_cosine(
+        bencher,
+        vectors(width, 0),
+        constant_vector(width),
+        NormMode::Exact,
+    );
 }
 
 /// One query vector represented as an extension array over constant storage.
@@ -124,5 +161,32 @@ fn extension_constant_vector(width: usize) -> ArrayRef {
 /// The rhs is the same broadcast query represented as extension-wrapped constant storage.
 #[divan::bench(args = WIDTHS)]
 fn column_x_extension_constant(bencher: Bencher, width: usize) {
-    bench_cosine(bencher, vectors(width, 0), extension_constant_vector(width));
+    bench_cosine(
+        bencher,
+        vectors(width, 0),
+        extension_constant_vector(width),
+        NormMode::Exact,
+    );
+}
+
+/// Measures exact cosine similarity over two [`Normalized`] inputs.
+#[divan::bench(args = WIDTHS)]
+fn normalized_exact(bencher: Bencher, width: usize) {
+    bench_cosine(
+        bencher,
+        normalized_vectors(width, 0),
+        normalized_vectors(width, 1),
+        NormMode::Exact,
+    );
+}
+
+/// Measures cosine similarity while trusting both normalized-direction claims.
+#[divan::bench(args = WIDTHS)]
+fn normalized_assume(bencher: Bencher, width: usize) {
+    bench_cosine(
+        bencher,
+        normalized_vectors(width, 0),
+        normalized_vectors(width, 1),
+        NormMode::AssumeNormalized,
+    );
 }

vortex-tensor/benches/l2_norm.rs

@@ -22,6 +22,8 @@ use vortex_array::arrays::MaskedArray;
 use vortex_array::arrays::PrimitiveArray;
 use vortex_array::validity::Validity;
 use vortex_buffer::Buffer;
+use vortex_tensor::encodings::normalized::Normalized;
+use vortex_tensor::scalar_fns::NormMode;
 use vortex_tensor::scalar_fns::l2_norm::L2Norm;
 use vortex_tensor::vector::Vector;
 
@@ -54,13 +56,32 @@ fn vectors(width: usize) -> ArrayRef {
     Vector::try_new_vector_array(storage).unwrap()
 }
 
-fn bench_l2_norm(bencher: Bencher, input: ArrayRef) {
+fn normalized_vectors(width: usize) -> ArrayRef {
+    let row_count = ELEMENTS / width;
+    let elements: Buffer<f64> = (0..ELEMENTS)
+        .map(|i| if i % width == 0 { 1.0 } else { 0.0 })
+        .collect();
+    let storage = FixedSizeListArray::new(
+        elements.into_array(),
+        u32::try_from(width).unwrap(),
+        Validity::NonNullable,
+        row_count,
+    )
+    .into_array();
+    let direction = Vector::try_new_vector_array(storage).unwrap();
+    let norms = PrimitiveArray::from_iter((1..=row_count).map(|norm| norm as f64)).into_array();
+
+    // SAFETY: Every direction row is unit length, and the non-negative norms have matching length.
+    unsafe { Normalized::new_unchecked(direction, norms, Validity::NonNullable) }.into_array()
+}
+
+fn bench_l2_norm(bencher: Bencher, input: ArrayRef, mode: NormMode) {
     let session = vortex_array::array_session();
     bencher
         .counter(ItemsCount::new(input.len()))
         .with_inputs(|| {
             (
-                L2Norm::try_new(input.clone()).unwrap().into_array(),
+                L2Norm::try_new(input.clone(), mode).unwrap().into_array(),
                 session.create_execution_ctx(),
             )
         })
@@ -69,7 +90,7 @@ fn bench_l2_norm(bencher: Bencher, input: ArrayRef) {
 
 #[divan::bench(args = WIDTHS)]
 fn non_nullable(bencher: Bencher, width: usize) {
-    bench_l2_norm(bencher, vectors(width));
+    bench_l2_norm(bencher, vectors(width), NormMode::Exact);
 }
 
 #[divan::bench(args = WIDTHS)]
@@ -78,5 +99,21 @@ fn nullable(bencher: Bencher, width: usize) {
     let input = MaskedArray::try_new(vectors(width), validity)
         .unwrap()
         .into_array();
-    bench_l2_norm(bencher, input);
+    bench_l2_norm(bencher, input, NormMode::Exact);
+}
+
+/// Measures the physical norm of a [`Normalized`] input.
+#[divan::bench(args = WIDTHS)]
+fn normalized_exact(bencher: Bencher, width: usize) {
+    bench_l2_norm(bencher, normalized_vectors(width), NormMode::Exact);
+}
+
+/// Reads the stored norm while trusting the normalized-direction claim.
+#[divan::bench(args = WIDTHS)]
+fn normalized_assume(bencher: Bencher, width: usize) {
+    bench_l2_norm(
+        bencher,
+        normalized_vectors(width),
+        NormMode::AssumeNormalized,
+    );
 }

vortex-tensor/src/encodings/normalized/array.rs

@@ -45,6 +45,10 @@ pub type NormalizedArray = Array<Normalized>;
 /// Row `i` decodes to `normalized[i] * norms[i]`. The encoding supports [`Vector`] and
 /// [`FixedShapeTensor`] columns with float elements.
 ///
+/// `Normalized` is a physical representation of those logical tensor types. The direction is
+/// evidence that norm-based scalar functions can use when their [`NormMode`] permits an
+/// approximation; it is not a logical refinement of [`Vector`] or [`FixedShapeTensor`].
+///
 /// # Invariants
 ///
 /// Every [`NormalizedArray`] has three slots.
@@ -56,7 +60,7 @@ pub type NormalizedArray = Array<Normalized>;
 /// Both data children have the array's length and element ptype. A missing validity slot means
 /// either non-nullable or nullable-all-valid data, as determined by the parent dtype.
 ///
-/// [`try_new`](Self::try_new) also enforces the invariants that make the split lossless:
+/// [`try_new`](Self::try_new) also validates the direction-and-norm relationship:
 ///
 /// - Each normalized row has L2 norm `1.0` or `0.0`, within the tolerance for its precision and
 ///   width.
@@ -68,16 +72,18 @@ pub type NormalizedArray = Array<Normalized>;
 /// # Lossy normalized children
 ///
 /// [`new_unchecked`](Self::new_unchecked) permits an approximate normalized child, such as a
-/// quantized direction. The stored norms remain authoritative. [`L2Norm`], [`InnerProduct`], and
-/// [`CosineSimilarity`] therefore operate on the stored children and can differ slightly from
-/// decoding and recomputing.
+/// quantized direction. [`NormMode::Exact`](crate::scalar_fns::NormMode::Exact) measures that
+/// physical direction. Only
+/// [`NormMode::AssumeNormalized`](crate::scalar_fns::NormMode::AssumeNormalized) trusts its norm as
+/// one, which can differ from decoding and recomputing and does not carry an error bound.
 ///
 /// [`Vector`]: crate::vector::Vector
 /// [`FixedShapeTensor`]: crate::fixed_shape_tensor::FixedShapeTensor
 /// [`normalize`]: crate::encodings::normalized::normalize
 /// [`L2Norm`]: crate::scalar_fns::l2_norm::L2Norm
 /// [`InnerProduct`]: crate::scalar_fns::inner_product::InnerProduct
 /// [`CosineSimilarity`]: crate::scalar_fns::cosine_similarity::CosineSimilarity
+/// [`NormMode`]: crate::scalar_fns::NormMode
 #[derive(Clone, Debug)]
 pub struct Normalized;
 

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

@@ -41,6 +41,7 @@ use crate::encodings::normalized::NormalizedArraySlotsExt;
 use crate::encodings::normalized::NormalizedSlots;
 use crate::encodings::normalized::array::DATA_CHILDREN;
 use crate::matcher::AnyTensor;
+use crate::scalar_fns::NormMode;
 use crate::scalar_fns::l2_norm::L2Norm;
 use crate::utils::extract_constant_flat_row;
 use crate::utils::extract_flat_elements;
@@ -137,7 +138,9 @@ pub fn normalize(input: ArrayRef, ctx: &mut ExecutionCtx) -> VortexResult<Normal
         return Ok(wrapped);
     }
 
-    let norms_array: ArrayRef = L2Norm::try_new(input.clone())?.into_array().execute(ctx)?;
+    let norms_array: ArrayRef = L2Norm::try_new(input.clone(), NormMode::Exact)?
+        .into_array()
+        .execute(ctx)?;
 
     // Execute before reading validity. Reading validity from the lazy array would run `L2Norm`
     // again when the values are requested.

vortex-tensor/src/encodings/normalized/mod.rs

@@ -5,12 +5,14 @@
 //!
 //! [`Normalized`] defines the physical layout and its invariants. Use [`normalize`] to create an
 //! exact split. [`L2Norm`], [`InnerProduct`], and [`CosineSimilarity`] can operate on the split
-//! without decoding it first.
+//! without decoding it first. Norm-based functions use [`NormMode`] to decide whether they measure
+//! the stored direction or trust its normalized-direction claim.
 //!
 //! [`normalize`]: crate::encodings::normalized::normalize
 //! [`L2Norm`]: crate::scalar_fns::l2_norm::L2Norm
 //! [`InnerProduct`]: crate::scalar_fns::inner_product::InnerProduct
 //! [`CosineSimilarity`]: crate::scalar_fns::cosine_similarity::CosineSimilarity
+//! [`NormMode`]: crate::scalar_fns::NormMode
 
 mod array;
 pub use array::Normalized;