usearch 2.21.3 → 2.25.0

package/numkong/rust/numkong.rs

@@ -0,0 +1,409 @@
+//! # NumKong - Hardware-Accelerated Numerics
+//!
+//! Provides SIMD-accelerated distance metrics, elementwise operations, and tensor algebra
+//! targeting ARM NEON/SVE/SME and x86 AVX2/AVX-512 backends.
+//!
+//! ## Modules
+//!
+//! - [`types`]: Mixed-precision scalar types (`f16`, `bf16`, FP8, packed integers) and [`FloatLike`] trait
+//! - [`spatial`]: Dot products, angular (cosine), and Euclidean distances
+//! - [`each`]: Elementwise operations and trigonometry
+//! - [`reduce`]: Statistical reductions (moments, min/max)
+//! - [`set`]: Binary set similarity (Hamming, Jaccard)
+//! - [`probability`]: Probability divergences (KL, JS)
+//! - [`curved`]: Curved metric spaces (Bilinear, Mahalanobis)
+//! - [`mesh`]: Mesh alignment (Kabsch, Umeyama, RMSD)
+//! - [`geospatial`]: Geospatial distances (Haversine, Vincenty)
+//! - [`sparse`]: Sparse set operations
+//! - [`cast`]: Type casting between scalar formats
+//! - [`capabilities`]: Runtime SIMD feature detection
+//! - [`matrix`]: Batch matrix operations (GEMM, packed spatial distances)
+//! - [`tensor`]: N-dimensional tensors with elementwise/reduction operations
+//!
+//! ## Implemented operations include:
+//!
+//! * Euclidean (L2), inner product, and angular (cosine) spatial distances.
+//! * Hamming and Jaccard binary distances.
+//! * Kullback-Leibler divergence and Jensen-Shannon distance.
+//! * Elementwise scale, sum, blend, and FMA operations.
+//! * Trigonometric functions (sin, cos, atan).
+//! * Type casting between all scalar formats.
+//! * Matrix multiplication with pre-packing (GEMM).
+//!
+//! ## Example
+//!
+//! ```rust
+//! use numkong::{Dot, Angular, Euclidean};
+//!
+//! let a = &[1.0_f32, 2.0, 3.0];
+//! let b = &[4.0_f32, 5.0, 6.0];
+//!
+//! let dot_product = f32::dot(a, b);
+//! let angular_dist = f32::angular(a, b);
+//! let l2sq_dist = f32::sqeuclidean(a, b);
+//!
+//! // Enable AMX and other platform-specific SIMD features
+//! numkong::capabilities::configure_thread();
+//! ```
+//!
+//! ## Mixed Precision Support
+//!
+//! ```rust
+//! use numkong::{Angular, f16, bf16};
+//!
+//! // Work with half-precision floats
+//! let half_a: Vec<f16> = vec![1.0, 2.0, 3.0].iter().map(|&x| f16::from_f32(x)).collect();
+//! let half_b: Vec<f16> = vec![4.0, 5.0, 6.0].iter().map(|&x| f16::from_f32(x)).collect();
+//! let half_angular_dist = f16::angular(&half_a, &half_b);
+//!
+//! // Work with brain floats
+//! let brain_a: Vec<bf16> = vec![1.0, 2.0, 3.0].iter().map(|&x| bf16::from_f32(x)).collect();
+//! let brain_b: Vec<bf16> = vec![4.0, 5.0, 6.0].iter().map(|&x| bf16::from_f32(x)).collect();
+//! let brain_angular_dist = bf16::angular(&brain_a, &brain_b);
+//!
+//! // Direct bit manipulation
+//! let half = f16::from_f32(3.14);
+//! let bits = half.0; // Access raw u16 representation
+//! let reconstructed = f16(bits);
+//! ```
+//!
+//! ## Traits
+//!
+//! The `SpatialSimilarity` trait (combining `Dot`, `Angular`, `Euclidean`) covers:
+//!
+//! - `dot(a, b)`: Computes dot product between two slices.
+//! - `angular(a, b)` / `cosine(a, b)`: Computes angular distance (1 − cosine similarity).
+//! - `sqeuclidean(a, b)`: Computes squared Euclidean distance.
+//! - `euclidean(a, b)`: Computes Euclidean distance.
+//!
+//! The `BinarySimilarity` trait (combining `Hamming`, `Jaccard`) covers:
+//!
+//! - `hamming(a, b)`: Computes Hamming distance between two slices.
+//! - `jaccard(a, b)`: Computes Jaccard distance between two slices.
+//!
+//! The `ProbabilitySimilarity` trait (combining `KullbackLeibler`, `JensenShannon`) covers:
+//!
+//! - `jensenshannon(a, b)`: Computes Jensen-Shannon distance.
+//! - `kullbackleibler(a, b)`: Computes Kullback-Leibler divergence.
+//!
+//! The elementwise traits (including `EachScale`, `EachSum`, `EachBlend`, `EachFMA`) covers:
+//!
+//! - `scale(a, alpha, beta, result)`: Element-wise `result[i] = α × a[i] + β`.
+//! - `sum(a, b, result)`: Element-wise `result[i] = a[i] + b[i]`.
+//! - `blend(a, b, alpha, beta, result)`: Blend `result[i] = α × a[i] + β × b[i]`.
+//! - `fma(a, b, c, alpha, beta, result)`: Fused multiply-add `result[i] = α × a[i] × b[i] + β × c[i]`.
+//!
+//! The `Trigonometry` trait (combining `EachSin`, `EachCos`, `EachATan`) covers:
+//!
+//! - `sin(input, result)`: Element-wise sine.
+//! - `cos(input, result)`: Element-wise cosine.
+//! - `atan(input, result)`: Element-wise arctangent.
+//!
+//! Additional traits: `VDot`, `Roots`, `SparseIntersect`, `SparseDot`.
+//!
+#![allow(non_camel_case_types)]
+#![allow(clippy::too_many_arguments)]
+#![cfg_attr(all(not(test), not(feature = "std")), no_std)]
+
+// Domain modules
+pub mod capabilities;
+pub mod cast;
+pub mod curved;
+pub mod each;
+pub mod geospatial;
+pub mod maxsim;
+pub mod mesh;
+pub mod probability;
+pub mod reduce;
+pub mod set;
+pub mod sparse;
+pub mod spatial;
+
+// Containers
+pub mod matrix;
+pub mod tensor;
+pub mod types;
+pub mod vector;
+
+// Re-export scalar types at crate root
+pub use types::{
+    bf16, bf16c, e2m3, e3m2, e4m3, e5m2, f16, f16c, f32c, f64c, i4x2, is_close, u1x8, u4x2, DimMut,
+    DimRef, FloatConvertible, FloatLike, NumberLike, StorageElement,
+};
+
+// Re-export spatial traits
+pub use spatial::{Angular, Dot, Euclidean, Roots, SpatialSimilarity, VDot};
+
+// Re-export set traits
+pub use set::{BinarySimilarity, Hamming, Jaccard};
+
+// Re-export probability traits
+pub use probability::{JensenShannon, KullbackLeibler, ProbabilitySimilarity};
+
+// Re-export elementwise and trig traits
+pub use each::{EachATan, EachBlend, EachCos, EachFMA, EachScale, EachSin, EachSum, Trigonometry};
+
+// Re-export reduction traits
+pub use reduce::{ReduceMinMax, ReduceMoments, Reductions};
+
+// Re-export curved metric traits
+pub use curved::{Bilinear, Mahalanobis};
+
+// Re-export mesh alignment
+pub use mesh::{MeshAlignment, MeshAlignmentResult};
+
+// Re-export geospatial
+pub use geospatial::{Geospatial, Haversine, Vincenty};
+
+// Re-export sparse
+pub use sparse::{SparseDot, SparseIntersect};
+
+// Re-export cast operations
+pub use cast::{cast, CastDtype};
+
+// Re-export capabilities
+pub use capabilities::cap;
+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,
+    SIMD_ALIGNMENT,
+};
+
+// Re-export matrix types
+pub use matrix::{
+    Angulars, Dots, Euclideans, Hammings, Jaccards, PackedMatrix, SymmetricAngulars, SymmetricDots,
+    SymmetricEuclideans, SymmetricHammings, SymmetricJaccards,
+};
+
+// Re-export vector types
+pub use vector::{
+    Vector, VectorIndex, VectorIterator, VectorSpan, VectorSpanIterator, VectorView,
+    VectorViewIterator,
+};
+
+// Re-export maxsim types
+pub use maxsim::{MaxSim, MaxSimPackedMatrix};
+
+// region: Tests
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn dot_smoke() {
+        let first = [1.0_f32, 2.0, 3.0];
+        let second = [4.0_f32, 5.0, 6.0];
+        assert!((<f32 as Dot>::dot(&first, &second).unwrap() - 32.0).abs() < 0.01);
+    }
+
+    #[test]
+    fn angular_smoke() {
+        let first = [1.0_f32, 0.0];
+        let second = [0.0_f32, 1.0];
+        // Orthogonal vectors → angular distance = 1.0
+        assert!((f32::angular(&first, &second).unwrap() - 1.0).abs() < 0.01);
+    }
+
+    #[test]
+    fn euclidean_smoke() {
+        let first = [0.0_f32, 0.0, 0.0];
+        let second = [3.0_f32, 4.0, 0.0];
+        assert!((f32::euclidean(&first, &second).unwrap() - 5.0).abs() < 0.1);
+    }
+
+    #[test]
+    fn maxsim_smoke() {
+        capabilities::configure_thread();
+        let queries = Tensor::<f32>::try_full(&[4, 16], 1.0).unwrap();
+        let documents = Tensor::<f32>::try_full(&[8, 16], 1.0).unwrap();
+        let queries_view = queries.view();
+        let docs_view = documents.view();
+        let queries_packed = MaxSimPackedMatrix::try_pack(&queries_view).unwrap();
+        let docs_packed = MaxSimPackedMatrix::try_pack(&docs_view).unwrap();
+        assert_eq!(queries_packed.dims(), (4, 16));
+        assert_eq!(docs_packed.dims(), (8, 16));
+        let score = queries_packed.score(&docs_packed);
+        assert!(
+            score.is_finite(),
+            "MaxSim score must be finite, got {score}"
+        );
+    }
+
+    #[test]
+    fn tensor_dots_smoke() {
+        capabilities::configure_thread();
+        let queries = Tensor::<f32>::try_full(&[2, 4], 1.0).unwrap();
+        let targets = Tensor::<f32>::try_full(&[3, 4], 1.0).unwrap();
+        let packed_targets = PackedMatrix::try_pack(&targets).unwrap();
+        let products = queries.dots_packed(&packed_targets);
+        assert_eq!(products.shape(), &[2, 3]);
+        assert!((products.as_slice()[0] - 4.0).abs() < 0.01);
+    }
+}
+
+// endregion: Tests
+
+// region: WASM Runtime Tests
+
+/// WASM runtime integration tests using Wasmtime
+/// These tests validate that WASI builds work correctly with standalone runtimes
+#[cfg(all(test, feature = "wasm-runtime"))]
+mod wasm_runtime_tests {
+    use std::fs;
+    use std::path::Path;
+    use wasmtime::{
+        Config, Engine, Extern, ExternType, Linker, Memory, MemoryType, Module, SharedMemory, Store,
+    };
+    use wasmtime_wasi::WasiCtx;
+
+    fn resolve_wasi_module() -> Option<String> {
+        if let Ok(path) = std::env::var("NK_WASI_MODULE") {
+            if Path::new(&path).exists() {
+                return Some(path);
+            }
+        }
+        if Path::new("build-wasi/nk_test.wasm").exists() {
+            Some("build-wasi/nk_test.wasm".to_string())
+        } else if Path::new("build-wasi/test.wasm").exists() {
+            Some("build-wasi/test.wasm".to_string())
+        } else {
+            None
+        }
+    }
+
+    /// Test that WASI WASM module can be loaded and executed with Wasmtime
+    /// This validates the dual-path capability detection (EM_ASM vs WASI imports)
+    #[test]
+    fn wasi_with_wasmtime() -> wasmtime::Result<()> {
+        // Check if WASI build exists
+        let Some(wasm_path) = resolve_wasi_module() else {
+            eprintln!("WASI build not found. Run:");
+            eprintln!("  export WASI_SDK_PATH=~/wasi-sdk");
+            eprintln!("  cmake -B build-wasi -DCMAKE_TOOLCHAIN_FILE=cmake/toolchain-wasi.cmake -DNK_BUILD_TEST=ON");
+            eprintln!("  cmake --build build-wasi --target nk_test");
+            return Ok(()); // Skip test if build doesn't exist
+        };
+
+        println!("Loading WASI module from {}", wasm_path);
+
+        let mut config = Config::new();
+        config.wasm_simd(true);
+        config.wasm_relaxed_simd(true);
+        config.wasm_threads(true);
+        config.shared_memory(true);
+
+        // Create Wasmtime engine and linker
+        let engine = Engine::new(&config)?;
+        let mut linker = Linker::new(&engine);
+
+        // Create WASI context (Wasmtime 41+ API)
+        // Don't inherit_args() — cargo's test filter args would confuse the WASM test binary.
+        let wasi = WasiCtx::builder().inherit_stdio().inherit_env().build_p1();
+        let mut store = Store::new(&engine, wasi);
+
+        // Add WASI support (Wasmtime 41+ requires p1 module)
+        wasmtime_wasi::p1::add_to_linker_sync(&mut linker, |s| s)?;
+
+        // Provide capability detection imports (required for WASI build)
+        // These functions are called from nk_capabilities_v128relaxed_() in C code
+        linker.func_wrap("env", "nk_has_v128", || -> i32 {
+            // Return 1 (true) - assume SIMD128 is available in Wasmtime
+            println!("  nk_has_v128() called from WASM -> returning 1");
+            1
+        })?;
+
+        linker.func_wrap("env", "nk_has_relaxed", || -> i32 {
+            // Return 1 (true) - assume Relaxed SIMD is available in Wasmtime
+            println!("  nk_has_relaxed() called from WASM -> returning 1");
+            1
+        })?;
+
+        linker.func_wrap("wasi", "thread-spawn", |_start_arg: i32| -> i32 { -1 })?;
+
+        // Load WASM module
+        let wasm_bytes = fs::read(&wasm_path)?;
+        let module = Module::new(&engine, wasm_bytes)?;
+
+        for import in module.imports() {
+            if import.module() != "env" || import.name() != "memory" {
+                continue;
+            }
+
+            let ExternType::Memory(memory_ty) = import.ty() else {
+                continue;
+            };
+
+            let minimum = u32::try_from(memory_ty.minimum()).map_err(|_| {
+                wasmtime::Error::msg(format!(
+                    "memory minimum {} does not fit in u32",
+                    memory_ty.minimum()
+                ))
+            })?;
+            let maximum = memory_ty
+                .maximum()
+                .ok_or_else(|| wasmtime::Error::msg("shared memory import is missing a maximum"))?;
+            let maximum = u32::try_from(maximum).map_err(|_| {
+                wasmtime::Error::msg(format!("memory maximum {maximum} does not fit in u32"))
+            })?;
+
+            if memory_ty.is_shared() {
+                let memory = SharedMemory::new(&engine, MemoryType::shared(minimum, maximum))?;
+                linker.define(&store, "env", "memory", Extern::from(memory))?;
+            } else {
+                let memory = Memory::new(&mut store, MemoryType::new(minimum, Some(maximum)))?;
+                linker.define(&store, "env", "memory", Extern::from(memory))?;
+            }
+        }
+
+        // Instantiate module
+        println!("Instantiating WASM module...");
+        let instance = linker.instantiate(&mut store, &module)?;
+
+        // Get _start entry point (WASI convention; exit code comes via proc_exit trap)
+        let start = instance.get_typed_func::<(), ()>(&mut store, "_start")?;
+
+        // Run tests — _start calls exit(0) which triggers proc_exit(0).
+        // Wasmtime reports proc_exit as an I32Exit error; exit code 0 means success.
+        println!("Running WASM tests...");
+        match start.call(&mut store, ()) {
+            Ok(()) => {}
+            Err(e) => {
+                if let Some(exit) = e.downcast_ref::<wasmtime_wasi::I32Exit>() {
+                    assert_eq!(exit.0, 0, "WASI tests failed with exit code {}", exit.0);
+                } else {
+                    return Err(e);
+                }
+            }
+        }
+
+        println!("WASM tests completed successfully");
+        Ok(())
+    }
+
+    /// Test capability detection mechanism works in WASI environment
+    #[test]
+    fn capability_imports() -> wasmtime::Result<()> {
+        println!("Testing capability import mechanism...");
+
+        // Create minimal engine for import testing
+        let engine = Engine::default();
+        let mut linker = Linker::<()>::new(&engine);
+
+        // Test that we can define the required imports
+        linker.func_wrap("env", "nk_has_v128", || -> i32 { 1 })?;
+        linker.func_wrap("env", "nk_has_relaxed", || -> i32 { 0 })?;
+
+        println!("  ✓ Capability imports defined successfully");
+
+        Ok(())
+    }
+}
+
+// endregion: WASM Runtime Tests

package/numkong/rust/probability.rs

@@ -0,0 +1,272 @@
+//! Probability measures: Kullback-Leibler divergence and Jensen-Shannon distance.
+//!
+//! This module provides:
+//!
+//! - [`KullbackLeibler`]: KL divergence between two distributions
+//! - [`JensenShannon`]: Jensen-Shannon distance (symmetric metric)
+//! - [`ProbabilitySimilarity`]: Blanket trait combining `KullbackLeibler + JensenShannon`
+
+use crate::types::{bf16, f16};
+
+#[link(name = "numkong")]
+extern "C" {
+    fn nk_jsd_f16(a: *const u16, b: *const u16, c: usize, d: *mut f32);
+    fn nk_jsd_bf16(a: *const u16, b: *const u16, c: usize, d: *mut f32);
+    fn nk_jsd_f32(a: *const f32, b: *const f32, c: usize, d: *mut f64);
+    fn nk_jsd_f64(a: *const f64, b: *const f64, c: usize, d: *mut f64);
+
+    fn nk_kld_f16(a: *const u16, b: *const u16, c: usize, d: *mut f32);
+    fn nk_kld_bf16(a: *const u16, b: *const u16, c: usize, d: *mut f32);
+    fn nk_kld_f32(a: *const f32, b: *const f32, c: usize, d: *mut f64);
+    fn nk_kld_f64(a: *const f64, b: *const f64, c: usize, d: *mut f64);
+}
+
+// region: KullbackLeibler
+
+/// Computes the **Kullback-Leibler divergence** between two probability distributions.
+///
+/// D_KL(P‖Q) = ∑ᵢ pᵢ × ln(pᵢ / qᵢ)
+///
+/// Range: \[0, ∞). Not symmetric. Returns `None` if lengths differ.
+///
+/// Implemented for: `f64`, `f32`, `f16`, `bf16`.
+pub trait KullbackLeibler: Sized {
+    type Output;
+    fn kullbackleibler(a: &[Self], b: &[Self]) -> Option<Self::Output>;
+
+    /// Alias for `kullbackleibler`.
+    fn kl(a: &[Self], b: &[Self]) -> Option<Self::Output> {
+        Self::kullbackleibler(a, b)
+    }
+}
+
+impl KullbackLeibler for f64 {
+    type Output = f64;
+    fn kullbackleibler(a: &[Self], b: &[Self]) -> Option<Self::Output> {
+        if a.len() != b.len() {
+            return None;
+        }
+        let mut result: Self::Output = 0.0;
+        unsafe { nk_kld_f64(a.as_ptr(), b.as_ptr(), a.len(), &mut result) };
+        Some(result)
+    }
+}
+
+impl KullbackLeibler for f32 {
+    type Output = f64;
+    fn kullbackleibler(a: &[Self], b: &[Self]) -> Option<Self::Output> {
+        if a.len() != b.len() {
+            return None;
+        }
+        let mut result: Self::Output = 0.0;
+        unsafe { nk_kld_f32(a.as_ptr(), b.as_ptr(), a.len(), &mut result) };
+        Some(result)
+    }
+}
+
+impl KullbackLeibler for f16 {
+    type Output = f32;
+    fn kullbackleibler(a: &[Self], b: &[Self]) -> Option<Self::Output> {
+        if a.len() != b.len() {
+            return None;
+        }
+        let mut result: Self::Output = 0.0;
+        unsafe {
+            nk_kld_f16(
+                a.as_ptr() as *const u16,
+                b.as_ptr() as *const u16,
+                a.len(),
+                &mut result,
+            )
+        };
+        Some(result)
+    }
+}
+
+impl KullbackLeibler for bf16 {
+    type Output = f32;
+    fn kullbackleibler(a: &[Self], b: &[Self]) -> Option<Self::Output> {
+        if a.len() != b.len() {
+            return None;
+        }
+        let mut result: Self::Output = 0.0;
+        unsafe {
+            nk_kld_bf16(
+                a.as_ptr() as *const u16,
+                b.as_ptr() as *const u16,
+                a.len(),
+                &mut result,
+            )
+        };
+        Some(result)
+    }
+}
+
+// endregion: KullbackLeibler
+
+// region: JensenShannon
+
+/// Computes the **Jensen-Shannon distance** between two probability distributions.
+///
+/// d_JS(P, Q) = √(½(D_KL(P‖M) + D_KL(Q‖M))), where M = (P + Q) / 2
+///
+/// Range: \[0, √ln2\]. Symmetric. Returns `None` if lengths differ.
+///
+/// Implemented for: `f64`, `f32`, `f16`, `bf16`.
+pub trait JensenShannon: Sized {
+    type Output;
+    fn jensenshannon(a: &[Self], b: &[Self]) -> Option<Self::Output>;
+
+    /// Alias for `jensenshannon`.
+    fn js(a: &[Self], b: &[Self]) -> Option<Self::Output> {
+        Self::jensenshannon(a, b)
+    }
+}
+
+impl JensenShannon for f64 {
+    type Output = f64;
+    fn jensenshannon(a: &[Self], b: &[Self]) -> Option<Self::Output> {
+        if a.len() != b.len() {
+            return None;
+        }
+        let mut result: Self::Output = 0.0;
+        unsafe { nk_jsd_f64(a.as_ptr(), b.as_ptr(), a.len(), &mut result) };
+        Some(result)
+    }
+}
+
+impl JensenShannon for f32 {
+    type Output = f64;
+    fn jensenshannon(a: &[Self], b: &[Self]) -> Option<Self::Output> {
+        if a.len() != b.len() {
+            return None;
+        }
+        let mut result: Self::Output = 0.0;
+        unsafe { nk_jsd_f32(a.as_ptr(), b.as_ptr(), a.len(), &mut result) };
+        Some(result)
+    }
+}
+
+impl JensenShannon for f16 {
+    type Output = f32;
+    fn jensenshannon(a: &[Self], b: &[Self]) -> Option<Self::Output> {
+        if a.len() != b.len() {
+            return None;
+        }
+        let mut result: Self::Output = 0.0;
+        unsafe {
+            nk_jsd_f16(
+                a.as_ptr() as *const u16,
+                b.as_ptr() as *const u16,
+                a.len(),
+                &mut result,
+            )
+        };
+        Some(result)
+    }
+}
+
+impl JensenShannon for bf16 {
+    type Output = f32;
+    fn jensenshannon(a: &[Self], b: &[Self]) -> Option<Self::Output> {
+        if a.len() != b.len() {
+            return None;
+        }
+        let mut result: Self::Output = 0.0;
+        unsafe {
+            nk_jsd_bf16(
+                a.as_ptr() as *const u16,
+                b.as_ptr() as *const u16,
+                a.len(),
+                &mut result,
+            )
+        };
+        Some(result)
+    }
+}
+
+// endregion: JensenShannon
+
+/// `ProbabilitySimilarity` bundles probability divergence metrics: KullbackLeibler and JensenShannon.
+pub trait ProbabilitySimilarity: KullbackLeibler + JensenShannon {}
+impl<Scalar: KullbackLeibler + JensenShannon> ProbabilitySimilarity for Scalar {}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use crate::types::{assert_close, bf16, f16, FloatLike, NumberLike, TestableType};
+
+    fn check_kld<Scalar>(distribution_a: &[f32], distribution_b: &[f32], expected: f64)
+    where
+        Scalar: FloatLike + TestableType + KullbackLeibler,
+        Scalar::Output: FloatLike,
+    {
+        let a_samples: Vec<Scalar> = distribution_a
+            .iter()
+            .map(|&v| Scalar::from_f32(v))
+            .collect();
+        let b_samples: Vec<Scalar> = distribution_b
+            .iter()
+            .map(|&v| Scalar::from_f32(v))
+            .collect();
+        let result = Scalar::kullbackleibler(&a_samples, &b_samples)
+            .unwrap()
+            .to_f64();
+        // Divergences involve ln() so need wider tolerance than simple dot products
+        assert_close(
+            result,
+            expected,
+            Scalar::atol().max(1e-6),
+            Scalar::rtol().max(1e-6),
+            &format!("kld<{}>", core::any::type_name::<Scalar>()),
+        );
+    }
+
+    fn check_jsd<Scalar>(distribution_a: &[f32], distribution_b: &[f32], expected: f64)
+    where
+        Scalar: FloatLike + TestableType + JensenShannon,
+        Scalar::Output: FloatLike,
+    {
+        let a_samples: Vec<Scalar> = distribution_a
+            .iter()
+            .map(|&v| Scalar::from_f32(v))
+            .collect();
+        let b_samples: Vec<Scalar> = distribution_b
+            .iter()
+            .map(|&v| Scalar::from_f32(v))
+            .collect();
+        let result = Scalar::jensenshannon(&a_samples, &b_samples)
+            .unwrap()
+            .to_f64();
+        // Divergences involve ln() so need wider tolerance than simple dot products
+        assert_close(
+            result,
+            expected,
+            Scalar::atol().max(1e-6),
+            Scalar::rtol().max(1e-6),
+            &format!("jsd<{}>", core::any::type_name::<Scalar>()),
+        );
+    }
+
+    #[test]
+    fn divergences() {
+        let distribution_a = &[0.1_f32, 0.9, 0.0];
+        let distribution_b = &[0.2_f32, 0.8, 0.0];
+
+        // KL(a||b) = 0.1*ln(0.1/0.2) + 0.9*ln(0.9/0.8)
+        let kld_expected = 0.1_f64 * (0.1_f64 / 0.2).ln() + 0.9_f64 * (0.9_f64 / 0.8).ln();
+        check_kld::<f64>(distribution_a, distribution_b, kld_expected);
+        check_kld::<f32>(distribution_a, distribution_b, kld_expected);
+        check_kld::<f16>(distribution_a, distribution_b, kld_expected);
+        check_kld::<bf16>(distribution_a, distribution_b, kld_expected);
+
+        // JS distance = sqrt(0.5 * (KL(a||m) + KL(b||m))) where m = (a+b)/2
+        let kl_am = 0.1_f64 * (0.1_f64 / 0.15).ln() + 0.9 * (0.9_f64 / 0.85).ln();
+        let kl_bm = 0.2_f64 * (0.2_f64 / 0.15).ln() + 0.8 * (0.8_f64 / 0.85).ln();
+        let jsd_expected = (0.5 * (kl_am + kl_bm)).sqrt();
+        check_jsd::<f64>(distribution_a, distribution_b, jsd_expected);
+        check_jsd::<f32>(distribution_a, distribution_b, jsd_expected);
+        check_jsd::<f16>(distribution_a, distribution_b, jsd_expected);
+        check_jsd::<bf16>(distribution_a, distribution_b, jsd_expected);
+    }
+}