usearch 2.26.1 → 2.26.2

package/numkong/rust/numkong.rs

@@ -141,10 +141,15 @@ pub use set::{BinarySimilarity, Hamming, Jaccard};
 pub use probability::{JensenShannon, KullbackLeibler, ProbabilitySimilarity};
 
 // Re-export elementwise and trig traits
-pub use each::{EachATan, EachBlend, EachCos, EachFMA, EachScale, EachSin, EachSum, Trigonometry};
+pub use each::{
+    AllCloseOps, BlendOps, EachATan, EachBlend, EachCos, EachFMA, EachScale, EachSin, EachSum,
+    FmaOps, ScaleOps, SumOps, TrigAtanOps, TrigCosOps, TrigSinOps, Trigonometry,
+};
 
 // Re-export reduction traits
-pub use reduce::{ReduceMinMax, ReduceMoments, Reductions};
+pub use reduce::{
+    BitwiseReductions, MinMaxOps, MomentsOps, ReduceMinMax, ReduceMoments, Reductions, SumSqToF64,
+};
 
 // Re-export curved metric traits
 pub use curved::{Bilinear, Mahalanobis};
@@ -159,7 +164,7 @@ pub use geospatial::{Geospatial, Haversine, Vincenty};
 pub use sparse::{SparseDot, SparseIntersect};
 
 // Re-export cast operations
-pub use cast::{cast, CastDtype};
+pub use cast::{cast, CastDtype, CastOps};
 
 // Re-export capabilities
 pub use capabilities::cap;
@@ -167,11 +172,10 @@ pub use capabilities::{available, configure_thread, uses_dynamic_dispatch};
 
 // Re-export tensor types
 pub use tensor::{
-    AllCloseOps, Allocator, AxisIterator, AxisIteratorMut, BlendOps, CastOps, FmaOps, Global,
-    Matrix, MatrixSpan, MatrixView, MinMaxOps, MinMaxResult, MomentsOps, RangeStep, ScaleOps,
-    SliceArg, SliceRange, SliceSpec, SumOps, Tensor, TensorDims, TensorError, TensorIterator,
-    TensorMut, TensorRef, TensorSpan, TensorSpanDims, TensorSpanIterator, TensorView,
-    TensorViewDims, TensorViewIterator, TrigAtanOps, TrigCosOps, TrigSinOps, DEFAULT_MAX_RANK,
+    Allocator, AxisIterator, AxisIteratorMut, CopyFrom, Fill, Global, Matrix, MatrixSpan,
+    MatrixView, MinMaxResult, RangeStep, SliceArg, SliceRange, SliceSpec, Tensor, TensorDims,
+    TensorError, TensorIterator, TensorMut, TensorRef, TensorSpan, TensorSpanDims,
+    TensorSpanIterator, TensorView, TensorViewDims, TensorViewIterator, DEFAULT_MAX_RANK,
     SIMD_ALIGNMENT,
 };
 

package/numkong/rust/reduce.rs

@@ -1,12 +1,24 @@
-//! Statistical reductions: moments (sum/sum-of-squares) and min/max.
+//! Statistical reductions: moments, min/max, and bit-population counts.
 //!
 //! This module provides:
 //!
-//! - [`ReduceMoments`]: Sum and sum-of-squares over a strided slice
-//! - [`ReduceMinMax`]: Minimum and maximum over a strided slice
-//! - [`Reductions`]: Blanket trait combining `ReduceMoments + ReduceMinMax`
-
+//! - Slice-level traits:
+//!   - [`ReduceMoments`]: Sum and sum-of-squares over a strided slice
+//!   - [`ReduceMinMax`]: Minimum and maximum over a strided slice
+//!   - [`Reductions`]: Blanket trait combining `ReduceMoments + ReduceMinMax`
+//! - Tensor-shaped extension traits (auto-implemented on every [`crate::tensor::TensorRef`]):
+//!   - [`MomentsOps`]: Per-axis and full-tensor moments / sums / norms
+//!   - [`MinMaxOps`]: Per-axis and full-tensor min / max / argmin / argmax
+//!   - [`BitwiseReductions`]: `popcount` / `any_set` / `none_set` / `all_set` for `u1x8` tensors;
+//!     `Vector<u1x8>` / `VectorView<u1x8>` / `VectorSpan<u1x8>` expose the same method names as
+//!     concrete impls (defined in [`crate::vector`]) until a `VectorRef` trait exists
+//! - [`SumSqToF64`]: Helper bound used by `MomentsOps::try_norm_*` to convert sum-of-squares to f64
+
+use crate::tensor::{
+    Global, MinMaxAxisResult, MinMaxResult, MomentsAxisResult, Tensor, TensorError, TensorRef,
+};
 use crate::types::{bf16, e2m3, e3m2, e4m3, e5m2, f16, i4x2, u1x8, u4x2, StorageElement};
+use crate::vector::VectorIndex;
 
 #[link(name = "numkong")]
 extern "C" {
@@ -1046,6 +1058,284 @@ impl ReduceMinMax for u1x8 {
 pub trait Reductions: ReduceMoments + ReduceMinMax {}
 impl<Scalar: ReduceMoments + ReduceMinMax> Reductions for Scalar {}
 
+/// Walks a `u1x8` tensor view, accumulating set-bit counts via
+/// [`u1x8::reduce_moments`] on each rank-1 contiguous lane.
+///
+/// At rank ≤ 1 the view's logical bit count divides cleanly into a contiguous
+/// storage slice (sub-byte stride is always one byte per `u1x8`, so rank-1 is
+/// trivially contiguous). At higher ranks we recurse on the leading axis via
+/// the `TensorRef`-supplied structural accessors so that matrices and higher-rank
+/// tensors of bits work without flattening first. Non-contiguous outer strides
+/// fall out of the recursion naturally — only the leaf storage slice has to be
+/// dense, and sub-byte storage always is.
+fn popcount_via_tensor_ref<View, const MAX_RANK: usize>(view: &View) -> u64
+where
+    View: TensorRef<u1x8, MAX_RANK> + ?Sized,
+{
+    let dims_per_value = u1x8::dimensions_per_value();
+    let logical_count: usize = view.shape().iter().product();
+    if logical_count == 0 {
+        return 0;
+    }
+    if view.ndim() <= 1 {
+        let storage_count = logical_count / dims_per_value;
+        if storage_count == 0 {
+            return 0;
+        }
+        let storage_slice = unsafe { core::slice::from_raw_parts(view.as_ptr(), storage_count) };
+        let (sum, _sum_of_squares) =
+            u1x8::reduce_moments(storage_slice, core::mem::size_of::<u1x8>());
+        return sum;
+    }
+    let mut total: u64 = 0;
+    let inner_view = view.view();
+    if let Ok(axis_iter) = inner_view.axis_views(0usize) {
+        for sub_view in axis_iter {
+            total += popcount_via_tensor_ref::<_, MAX_RANK>(&sub_view);
+        }
+    }
+    total
+}
+
+/// Population count, any/none/all reductions over packed-bit (`u1x8`) tensors.
+///
+/// Blanket-impl'd on every [`TensorRef<u1x8, R>`], so `Tensor<u1x8>`,
+/// `TensorView<u1x8>`, and `TensorSpan<u1x8>` all expose `.popcount()` /
+/// `.any_set()` / `.none_set()` / `.all_set()` with identical spellings.
+/// Mirrors `numkong::popcount(tensor_view)` in `include/numkong/reduce.hpp`.
+///
+/// Method names use the `_set` suffix to leave room for a future generic
+/// `any` / `none` / `all` over arbitrary scalar predicates and to avoid
+/// colliding with the `all` slice marker exported from [`crate::vector`].
+pub trait BitwiseReductions<const MAX_RANK: usize>: TensorRef<u1x8, MAX_RANK> {
+    /// Number of set bits across the entire tensor.
+    fn popcount(&self) -> u64 {
+        popcount_via_tensor_ref::<_, MAX_RANK>(self)
+    }
+
+    /// `true` if at least one bit in the tensor is set.
+    fn any_set(&self) -> bool {
+        self.popcount() != 0
+    }
+
+    /// `true` if no bit in the tensor is set.
+    fn none_set(&self) -> bool {
+        !self.any_set()
+    }
+
+    /// `true` if every bit in the tensor is set.
+    fn all_set(&self) -> bool {
+        self.popcount() == self.numel() as u64
+    }
+}
+
+impl<Container, const MAX_RANK: usize> BitwiseReductions<MAX_RANK> for Container where
+    Container: TensorRef<u1x8, MAX_RANK> + ?Sized
+{
+}
+
+// region: SumSqToF64 helper trait (moved from crate::tensor)
+
+#[doc(hidden)]
+pub trait SumSqToF64 {
+    fn to_f64(self) -> f64;
+}
+
+impl SumSqToF64 for f32 {
+    fn to_f64(self) -> f64 {
+        self as f64
+    }
+}
+impl SumSqToF64 for f64 {
+    fn to_f64(self) -> f64 {
+        self
+    }
+}
+impl SumSqToF64 for u64 {
+    fn to_f64(self) -> f64 {
+        self as f64
+    }
+}
+impl SumSqToF64 for i64 {
+    fn to_f64(self) -> f64 {
+        self as f64
+    }
+}
+
+// endregion: SumSqToF64
+
+// region: MomentsOps / MinMaxOps (moved from crate::tensor)
+
+/// Extension trait: statistical reductions for any [`TensorRef`] implementor.
+pub trait MomentsOps<Scalar: Clone + ReduceMoments, const MAX_RANK: usize>:
+    TensorRef<Scalar, MAX_RANK>
+where
+    Scalar::SumOutput: Clone + Default + core::ops::AddAssign,
+    Scalar::SumSqOutput: Clone + Default + core::ops::AddAssign + SumSqToF64,
+{
+    fn try_moments_all(&self) -> Result<(Scalar::SumOutput, Scalar::SumSqOutput), TensorError> {
+        self.view().try_moments_all()
+    }
+
+    fn try_moments_axis<AnyIndex: VectorIndex>(
+        &self,
+        axis: AnyIndex,
+        keep_dims: bool,
+    ) -> MomentsAxisResult<Scalar, MAX_RANK> {
+        self.view().try_moments_axis(axis, keep_dims)
+    }
+
+    fn try_moments_axis_into<AnyIndex: VectorIndex>(
+        &self,
+        axis: AnyIndex,
+        keep_dims: bool,
+        sum_out: &mut Tensor<Scalar::SumOutput, Global, MAX_RANK>,
+        sumsq_out: &mut Tensor<Scalar::SumSqOutput, Global, MAX_RANK>,
+    ) -> Result<(), TensorError> {
+        self.view()
+            .try_moments_axis_into(axis, keep_dims, sum_out, sumsq_out)
+    }
+
+    fn try_sum_all(&self) -> Result<Scalar::SumOutput, TensorError> {
+        self.view().try_sum_all()
+    }
+
+    fn try_sum_axis<AnyIndex: VectorIndex>(
+        &self,
+        axis: AnyIndex,
+        keep_dims: bool,
+    ) -> Result<Tensor<Scalar::SumOutput, Global, MAX_RANK>, TensorError> {
+        self.view().try_sum_axis(axis, keep_dims)
+    }
+
+    fn try_sum_axis_into<AnyIndex: VectorIndex>(
+        &self,
+        axis: AnyIndex,
+        keep_dims: bool,
+        out: &mut Tensor<Scalar::SumOutput, Global, MAX_RANK>,
+    ) -> Result<(), TensorError> {
+        self.view().try_sum_axis_into(axis, keep_dims, out)
+    }
+
+    fn try_norm_all(&self) -> Result<f64, TensorError> {
+        self.view().try_norm_all()
+    }
+
+    fn try_norm_axis<AnyIndex: VectorIndex>(
+        &self,
+        axis: AnyIndex,
+        keep_dims: bool,
+    ) -> Result<Tensor<f64, Global, MAX_RANK>, TensorError> {
+        self.view().try_norm_axis(axis, keep_dims)
+    }
+
+    fn try_norm_axis_into<AnyIndex: VectorIndex>(
+        &self,
+        axis: AnyIndex,
+        keep_dims: bool,
+        out: &mut Tensor<f64, Global, MAX_RANK>,
+    ) -> Result<(), TensorError> {
+        self.view().try_norm_axis_into(axis, keep_dims, out)
+    }
+}
+
+impl<Scalar: Clone + ReduceMoments, const R: usize, C: TensorRef<Scalar, R>> MomentsOps<Scalar, R>
+    for C
+where
+    Scalar::SumOutput: Clone + Default + core::ops::AddAssign,
+    Scalar::SumSqOutput: Clone + Default + core::ops::AddAssign + SumSqToF64,
+{
+}
+
+/// Extension trait: min/max reductions for any [`TensorRef`] implementor.
+pub trait MinMaxOps<Scalar: Clone + ReduceMinMax, const MAX_RANK: usize>:
+    TensorRef<Scalar, MAX_RANK>
+where
+    Scalar::Output: Clone + Default + PartialOrd,
+{
+    fn try_minmax_all(&self) -> Result<MinMaxResult<Scalar::Output>, TensorError> {
+        self.view().try_minmax_all()
+    }
+
+    fn try_minmax_axis<AnyIndex: VectorIndex>(
+        &self,
+        axis: AnyIndex,
+        keep_dims: bool,
+    ) -> MinMaxAxisResult<Scalar, MAX_RANK> {
+        self.view().try_minmax_axis(axis, keep_dims)
+    }
+
+    fn try_minmax_axis_into<AnyIndex: VectorIndex>(
+        &self,
+        axis: AnyIndex,
+        keep_dims: bool,
+        min_out: &mut Tensor<Scalar::Output, Global, MAX_RANK>,
+        argmin_out: &mut Tensor<usize, Global, MAX_RANK>,
+        max_out: &mut Tensor<Scalar::Output, Global, MAX_RANK>,
+        argmax_out: &mut Tensor<usize, Global, MAX_RANK>,
+    ) -> Result<(), TensorError> {
+        self.view()
+            .try_minmax_axis_into(axis, keep_dims, min_out, argmin_out, max_out, argmax_out)
+    }
+
+    fn try_min_all(&self) -> Result<Scalar::Output, TensorError> {
+        self.view().try_min_all()
+    }
+
+    fn try_argmin_all(&self) -> Result<usize, TensorError> {
+        self.view().try_argmin_all()
+    }
+
+    fn try_max_all(&self) -> Result<Scalar::Output, TensorError> {
+        self.view().try_max_all()
+    }
+
+    fn try_argmax_all(&self) -> Result<usize, TensorError> {
+        self.view().try_argmax_all()
+    }
+
+    fn try_min_axis<AnyIndex: VectorIndex>(
+        &self,
+        axis: AnyIndex,
+        keep_dims: bool,
+    ) -> Result<Tensor<Scalar::Output, Global, MAX_RANK>, TensorError> {
+        self.view().try_min_axis(axis, keep_dims)
+    }
+
+    fn try_argmin_axis<AnyIndex: VectorIndex>(
+        &self,
+        axis: AnyIndex,
+        keep_dims: bool,
+    ) -> Result<Tensor<usize, Global, MAX_RANK>, TensorError> {
+        self.view().try_argmin_axis(axis, keep_dims)
+    }
+
+    fn try_max_axis<AnyIndex: VectorIndex>(
+        &self,
+        axis: AnyIndex,
+        keep_dims: bool,
+    ) -> Result<Tensor<Scalar::Output, Global, MAX_RANK>, TensorError> {
+        self.view().try_max_axis(axis, keep_dims)
+    }
+
+    fn try_argmax_axis<AnyIndex: VectorIndex>(
+        &self,
+        axis: AnyIndex,
+        keep_dims: bool,
+    ) -> Result<Tensor<usize, Global, MAX_RANK>, TensorError> {
+        self.view().try_argmax_axis(axis, keep_dims)
+    }
+}
+
+impl<Scalar: Clone + ReduceMinMax, const R: usize, C: TensorRef<Scalar, R>> MinMaxOps<Scalar, R>
+    for C
+where
+    Scalar::Output: Clone + Default + PartialOrd,
+{
+}
+
+// endregion: MomentsOps / MinMaxOps
+
 #[cfg(test)]
 mod tests {
     use super::*;
@@ -1246,4 +1536,142 @@ mod tests {
     }
 
     // endregion
+
+    // region: BitwiseReductions across containers
+
+    #[test]
+    fn bitwise_popcount_across_tensor_shapes() {
+        use crate::tensor::Tensor;
+        // 16 bytes × 8 bits = 128 bits total. Alternating fully-set/zero bytes → popcount = 64.
+        let bits_storage: Vec<u1x8> = (0..16u8)
+            .map(|byte_index| if byte_index % 2 == 0 { u1x8(0xFFu8) } else { u1x8(0x00u8) })
+            .collect();
+        let mut tensor = Tensor::<u1x8>::try_zeros(&[4, 32]).unwrap();
+        for (slot, value) in tensor.as_mut_slice().iter_mut().zip(bits_storage.iter()) {
+            *slot = *value;
+        }
+
+        // Same `.popcount()` spelling on Tensor and TensorView.
+        assert_eq!(tensor.popcount(), 64);
+        assert_eq!(tensor.view().popcount(), 64);
+
+        assert!(tensor.any_set());
+        assert!(tensor.view().any_set());
+
+        assert!(!tensor.none_set());
+        assert!(!tensor.view().none_set());
+
+        assert!(!tensor.all_set());
+        assert!(!tensor.view().all_set());
+
+        // All-zeros tensor → none_set, !any_set, !all_set.
+        let zero_tensor = Tensor::<u1x8>::try_zeros(&[4, 32]).unwrap();
+        assert_eq!(zero_tensor.popcount(), 0);
+        assert!(!zero_tensor.any_set());
+        assert!(zero_tensor.none_set());
+        assert!(!zero_tensor.all_set());
+
+        // All-ones tensor → all_set.
+        let mut ones_tensor = Tensor::<u1x8>::try_zeros(&[4, 32]).unwrap();
+        for slot in ones_tensor.as_mut_slice().iter_mut() {
+            *slot = u1x8(0xFFu8);
+        }
+        assert_eq!(ones_tensor.popcount(), 128);
+        assert!(ones_tensor.any_set());
+        assert!(!ones_tensor.none_set());
+        assert!(ones_tensor.all_set());
+    }
+
+    #[test]
+    fn bitwise_reductions_on_vector_containers() {
+        use crate::vector::{Vector, VectorSpan, VectorView};
+        // Three set bytes (24 bits), one zero byte → popcount = 24.
+        let mut storage = [u1x8(0xFFu8), u1x8(0xFFu8), u1x8(0x00u8), u1x8(0xFFu8)];
+        let mut vector = Vector::<u1x8>::try_zeros(32).unwrap();
+        for (slot, value) in vector.as_mut_slice().iter_mut().zip(storage.iter()) {
+            *slot = *value;
+        }
+        assert_eq!(vector.popcount(), 24);
+        assert!(vector.any_set());
+        assert!(!vector.none_set());
+        assert!(!vector.all_set());
+
+        let view = unsafe {
+            VectorView::<u1x8>::from_raw_parts(
+                storage.as_ptr(),
+                32,
+                core::mem::size_of::<u1x8>() as isize,
+            )
+        };
+        assert_eq!(view.popcount(), 24);
+        assert!(view.any_set());
+
+        let span = unsafe {
+            VectorSpan::<u1x8>::from_raw_parts(
+                storage.as_mut_ptr(),
+                32,
+                core::mem::size_of::<u1x8>() as isize,
+            )
+        };
+        assert_eq!(span.popcount(), 24);
+        assert!(span.any_set());
+    }
+
+    // endregion: BitwiseReductions across containers
+
+    // region: tensor-shaped MomentsOps / MinMaxOps wrappers
+
+    #[test]
+    fn reductions_axis_and_strided_views() {
+        use crate::tensor::{MinMaxResult, SliceRange, Tensor};
+        let data: Vec<f32> = (0..12).map(|i| i as f32).collect();
+        let a = Tensor::<f32>::try_from_slice(&data, &[3, 4]).unwrap();
+        let a_even = a
+            .slice(&[SliceRange::full(), SliceRange::range_step(0, 4, 2)])
+            .unwrap();
+
+        let sum_all = a_even.try_sum_all().unwrap();
+        assert!((sum_all - 30.0).abs() < 1e-6);
+
+        let norm_all = a_even.try_norm_all().unwrap();
+        assert!((norm_all - 14.832396974191326).abs() < 1e-9);
+
+        let (sum_axis0, sumsq_axis0) = a_even.try_moments_axis(0, false).unwrap();
+        assert_eq!(sum_axis0.shape(), &[2]);
+        assert!((sum_axis0.as_slice()[0] - 12.0).abs() < 1e-6);
+        assert!((sum_axis0.as_slice()[1] - 18.0).abs() < 1e-6);
+        assert!((sumsq_axis0.as_slice()[0] - 80.0).abs() < 1e-6);
+        assert!((sumsq_axis0.as_slice()[1] - 140.0).abs() < 1e-6);
+
+        let sum_axis1_keep = a_even.try_sum_axis(-1_i32, true).unwrap();
+        assert_eq!(sum_axis1_keep.shape(), &[3, 1]);
+        assert!((sum_axis1_keep.as_slice()[0] - 2.0).abs() < 1e-6);
+        assert!((sum_axis1_keep.as_slice()[1] - 10.0).abs() < 1e-6);
+        assert!((sum_axis1_keep.as_slice()[2] - 18.0).abs() < 1e-6);
+
+        let MinMaxResult {
+            min_value: min_axis0,
+            min_index: argmin_axis0,
+            max_value: max_axis0,
+            max_index: argmax_axis0,
+        } = a_even.try_minmax_axis(0, false).unwrap();
+        assert_eq!(min_axis0.as_slice(), &[0.0, 2.0]);
+        assert_eq!(max_axis0.as_slice(), &[8.0, 10.0]);
+        assert_eq!(argmin_axis0.as_slice(), &[0, 0]);
+        assert_eq!(argmax_axis0.as_slice(), &[2, 2]);
+
+        let reversed = a
+            .slice(&[SliceRange::full(), SliceRange::range_step(3, 0, -1)])
+            .unwrap();
+        let reversed_sum = reversed.try_sum_axis(-1_i32, false).unwrap();
+        assert_eq!(reversed_sum.shape(), &[3]);
+        assert_eq!(reversed_sum.as_slice(), &[6.0, 18.0, 30.0]);
+
+        let reversed_argmin = reversed.try_argmin_axis(-1_i32, false).unwrap();
+        let reversed_argmax = reversed.try_argmax_axis(-1_i32, false).unwrap();
+        assert_eq!(reversed_argmin.as_slice(), &[2, 2, 2]);
+        assert_eq!(reversed_argmax.as_slice(), &[0, 0, 0]);
+    }
+
+    // endregion: tensor-shaped MomentsOps / MinMaxOps wrappers
 }