numkong 7.0.0

package/include/numkong/probability.h

@@ -0,0 +1,383 @@
+/**
+ *  @brief SIMD-accelerated Similarity Measures for Probability Distributions.
+ *  @file include/numkong/probability.h
+ *  @author Ash Vardanian
+ *  @date October 20, 2023
+ *
+ *  Contains following similarity measures:
+ *
+ *  - Kullback-Leibler Divergence (KLD)
+ *  - Jensen-Shannon Distance (JSD)
+ *
+ *  For dtypes:
+ *
+ *  - 64-bit floating point numbers → 64-bit
+ *  - 32-bit floating point numbers → 64-bit
+ *  - 16-bit floating point numbers → 32-bit
+ *  - 16-bit brain-floating point numbers → 32-bit
+ *
+ *  Precision policy:
+ *
+ *  - For `f32` inputs, the per-element vertical path stays in `f32` to preserve the fast ratio/log
+ *    approximations and SIMD throughput.
+ *  - The horizontal reduction over those per-element contributions widens to `f64`, and public
+ *    `f32` results are exposed as `f64`.
+ *  - For `f64` inputs, both the vertical path and the horizontal reduction stay in `f64`, with
+ *    stable summation in the serial kernels.
+ *  - For `f16` and `bf16` inputs, the kernels still widen to `f32`.
+ *
+ *  For hardware architectures:
+ *
+ *  - Arm: NEON
+ *  - x86: Haswell, Skylake, Sapphire
+ *
+ *  @section x86_instructions Relevant x86 Instructions
+ *
+ *  KL/JS divergence requires log2(x) which decomposes into exponent extraction (VGETEXP) plus
+ *  mantissa polynomial (using VGETMANT + FMA chain). This approach is faster than scalar log()
+ *  calls. Division (for p/q ratio) uses either VDIVPS directly or VRCP14PS with Newton-Raphson
+ *  refinement when ~14-bit precision suffices. Genoa's VGETEXP/VGETMANT are 25% faster than Ice.
+ *
+ *      Intrinsic               Instruction                     Ice         Genoa
+ *      _mm512_getexp_ps        VGETEXPPS (ZMM, ZMM)            4c @ p0     3c @ p23
+ *      _mm512_getexp_pd        VGETEXPPD (ZMM, ZMM)            4c @ p0     3c @ p23
+ *      _mm512_getmant_ps       VGETMANTPS (ZMM, ZMM, I8)       4c @ p0     3c @ p23
+ *      _mm512_getmant_pd       VGETMANTPD (ZMM, ZMM, I8)       4c @ p0     3c @ p23
+ *      _mm512_rcp14_ps         VRCP14PS (ZMM, ZMM)             7c @ p05    5c @ p01
+ *      _mm512_div_ps           VDIVPS (ZMM, ZMM, ZMM)          17c @ p05   11c @ p01
+ *      _mm512_fmadd_ps         VFMADD231PS (ZMM, ZMM, ZMM)     4c @ p0     4c @ p01
+ *
+ *  @section arm_instructions Relevant ARM NEON/SVE Instructions
+ *
+ *  ARM lacks direct exponent/mantissa extraction, so log2 uses integer reinterpretation of the
+ *  float bits followed by polynomial refinement. FRECPE provides ~8-bit reciprocal approximation
+ *  for division, refined with FRECPS Newton-Raphson steps to ~22-bit precision.
+ *
+ *      Intrinsic               Instruction     M1 Firestorm    Graviton 3      Graviton 4
+ *      vfmaq_f32               FMLA.S (vec)    4c @ V0123      4c @ V0123      4c @ V0123
+ *      vrecpeq_f32             FRECPE.S        3c @ V02        3c @ V02        3c @ V02
+ *      vrecpsq_f32             FRECPS.S        4c @ V0123      4c @ V0123      4c @ V0123
+ *
+ *  @section references References
+ *
+ *  - x86 intrinsics: https://www.intel.com/content/www/us/en/docs/intrinsics-guide/
+ *  - Arm intrinsics: https://developer.arm.com/architectures/instruction-sets/intrinsics/
+ *
+ */
+#ifndef NK_PROBABILITY_H
+#define NK_PROBABILITY_H
+
+#include "numkong/types.h"
+#include "numkong/reduce.h" // For horizontal reduction helpers
+
+#if defined(__cplusplus)
+extern "C" {
+#endif
+
+/**
+ *  @brief Kullback-Leibler divergence between two discrete probability distributions.
+ *
+ *  @param[in] a The first discrete probability distribution.
+ *  @param[in] b The second discrete probability distribution.
+ *  @param[in] n The number of elements in the distributions.
+ *  @param[out] result The output divergence value.
+ *
+ *  @note The distributions are assumed to be normalized.
+ *  @note The output divergence value is non-negative.
+ *  @note The output divergence value is zero if and only if the two distributions are identical.
+ */
+NK_DYNAMIC void nk_kld_f16(nk_f16_t const *a, nk_f16_t const *b, nk_size_t n, nk_f32_t *result);
+/**
+ *  @brief Kullback-Leibler divergence between two discrete probability distributions.
+ *
+ *  @param[in] a The first discrete probability distribution.
+ *  @param[in] b The second discrete probability distribution.
+ *  @param[in] n The number of elements in the distributions.
+ *  @param[out] result The output divergence value.
+ *
+ *  @note The distributions are assumed to be normalized.
+ *  @note The output divergence value is non-negative.
+ *  @note The output divergence value is zero if and only if the two distributions are identical.
+ */
+NK_DYNAMIC void nk_kld_bf16(nk_bf16_t const *a, nk_bf16_t const *b, nk_size_t n, nk_f32_t *result);
+/**
+ *  @brief Kullback-Leibler divergence between two discrete probability distributions.
+ *
+ *  @param[in] a The first discrete probability distribution.
+ *  @param[in] b The second discrete probability distribution.
+ *  @param[in] n The number of elements in the distributions.
+ *  @param[out] result The output divergence value.
+ *
+ *  @note The distributions are assumed to be normalized.
+ *  @note The output divergence value is non-negative.
+ *  @note The output divergence value is zero if and only if the two distributions are identical.
+ */
+NK_DYNAMIC void nk_kld_f32(nk_f32_t const *a, nk_f32_t const *b, nk_size_t n, nk_f64_t *result);
+/**
+ *  @brief Kullback-Leibler divergence between two discrete probability distributions.
+ *
+ *  @param[in] a The first discrete probability distribution.
+ *  @param[in] b The second discrete probability distribution.
+ *  @param[in] n The number of elements in the distributions.
+ *  @param[out] result The output divergence value.
+ *
+ *  @note The distributions are assumed to be normalized.
+ *  @note The output divergence value is non-negative.
+ *  @note The output divergence value is zero if and only if the two distributions are identical.
+ */
+NK_DYNAMIC void nk_kld_f64(nk_f64_t const *a, nk_f64_t const *b, nk_size_t n, nk_f64_t *result);
+/**
+ *  @brief Jensen-Shannon distance between two discrete probability distributions.
+ *
+ *  @param[in] a The first discrete probability distribution.
+ *  @param[in] b The second discrete probability distribution.
+ *  @param[in] n The number of elements in the distributions.
+ *  @param[out] result The output distance value.
+ *
+ *  @note The distributions are assumed to be normalized.
+ *  @note The output distance value is non-negative.
+ *  @note The output distance value is zero if and only if the two distributions are identical.
+ */
+NK_DYNAMIC void nk_jsd_f16(nk_f16_t const *a, nk_f16_t const *b, nk_size_t n, nk_f32_t *result);
+/**
+ *  @brief Jensen-Shannon distance between two discrete probability distributions.
+ *
+ *  @param[in] a The first discrete probability distribution.
+ *  @param[in] b The second discrete probability distribution.
+ *  @param[in] n The number of elements in the distributions.
+ *  @param[out] result The output distance value.
+ *
+ *  @note The distributions are assumed to be normalized.
+ *  @note The output distance value is non-negative.
+ *  @note The output distance value is zero if and only if the two distributions are identical.
+ */
+NK_DYNAMIC void nk_jsd_bf16(nk_bf16_t const *a, nk_bf16_t const *b, nk_size_t n, nk_f32_t *result);
+/**
+ *  @brief Jensen-Shannon distance between two discrete probability distributions.
+ *
+ *  @param[in] a The first discrete probability distribution.
+ *  @param[in] b The second discrete probability distribution.
+ *  @param[in] n The number of elements in the distributions.
+ *  @param[out] result The output distance value.
+ *
+ *  @note The distributions are assumed to be normalized.
+ *  @note The output distance value is non-negative.
+ *  @note The output distance value is zero if and only if the two distributions are identical.
+ */
+NK_DYNAMIC void nk_jsd_f32(nk_f32_t const *a, nk_f32_t const *b, nk_size_t n, nk_f64_t *result);
+/**
+ *  @brief Jensen-Shannon distance between two discrete probability distributions.
+ *
+ *  @param[in] a The first discrete probability distribution.
+ *  @param[in] b The second discrete probability distribution.
+ *  @param[in] n The number of elements in the distributions.
+ *  @param[out] result The output distance value.
+ *
+ *  @note The distributions are assumed to be normalized.
+ *  @note The output distance value is non-negative.
+ *  @note The output distance value is zero if and only if the two distributions are identical.
+ */
+NK_DYNAMIC void nk_jsd_f64(nk_f64_t const *a, nk_f64_t const *b, nk_size_t n, nk_f64_t *result);
+
+/** @copydoc nk_kld_f64 */
+NK_PUBLIC void nk_kld_f64_serial(nk_f64_t const *a, nk_f64_t const *b, nk_size_t n, nk_f64_t *result);
+/** @copydoc nk_jsd_f64 */
+NK_PUBLIC void nk_jsd_f64_serial(nk_f64_t const *a, nk_f64_t const *b, nk_size_t n, nk_f64_t *result);
+/** @copydoc nk_kld_f32 */
+NK_PUBLIC void nk_kld_f32_serial(nk_f32_t const *a, nk_f32_t const *b, nk_size_t n, nk_f64_t *result);
+/** @copydoc nk_jsd_f32 */
+NK_PUBLIC void nk_jsd_f32_serial(nk_f32_t const *a, nk_f32_t const *b, nk_size_t n, nk_f64_t *result);
+/** @copydoc nk_kld_f16 */
+NK_PUBLIC void nk_kld_f16_serial(nk_f16_t const *a, nk_f16_t const *b, nk_size_t n, nk_f32_t *result);
+/** @copydoc nk_jsd_f16 */
+NK_PUBLIC void nk_jsd_f16_serial(nk_f16_t const *a, nk_f16_t const *b, nk_size_t n, nk_f32_t *result);
+/** @copydoc nk_kld_bf16 */
+NK_PUBLIC void nk_kld_bf16_serial(nk_bf16_t const *a, nk_bf16_t const *b, nk_size_t n, nk_f32_t *result);
+/** @copydoc nk_jsd_bf16 */
+NK_PUBLIC void nk_jsd_bf16_serial(nk_bf16_t const *a, nk_bf16_t const *b, nk_size_t n, nk_f32_t *result);
+
+#if NK_TARGET_NEON
+/** @copydoc nk_kld_f32 */
+NK_PUBLIC void nk_kld_f32_neon(nk_f32_t const *a, nk_f32_t const *b, nk_size_t n, nk_f64_t *result);
+/** @copydoc nk_jsd_f32 */
+NK_PUBLIC void nk_jsd_f32_neon(nk_f32_t const *a, nk_f32_t const *b, nk_size_t n, nk_f64_t *result);
+#endif // NK_TARGET_NEON
+
+#if NK_TARGET_NEONHALF
+/** @copydoc nk_kld_f16 */
+NK_PUBLIC void nk_kld_f16_neonhalf(nk_f16_t const *a, nk_f16_t const *b, nk_size_t n, nk_f32_t *result);
+/** @copydoc nk_jsd_f16 */
+NK_PUBLIC void nk_jsd_f16_neonhalf(nk_f16_t const *a, nk_f16_t const *b, nk_size_t n, nk_f32_t *result);
+#endif // NK_TARGET_NEONHALF
+
+#if NK_TARGET_HASWELL
+/** @copydoc nk_kld_f64 */
+NK_PUBLIC void nk_kld_f64_haswell(nk_f64_t const *a, nk_f64_t const *b, nk_size_t n, nk_f64_t *result);
+/** @copydoc nk_jsd_f64 */
+NK_PUBLIC void nk_jsd_f64_haswell(nk_f64_t const *a, nk_f64_t const *b, nk_size_t n, nk_f64_t *result);
+/** @copydoc nk_kld_f16 */
+NK_PUBLIC void nk_kld_f16_haswell(nk_f16_t const *a, nk_f16_t const *b, nk_size_t n, nk_f32_t *result);
+/** @copydoc nk_jsd_f16 */
+NK_PUBLIC void nk_jsd_f16_haswell(nk_f16_t const *a, nk_f16_t const *b, nk_size_t n, nk_f32_t *result);
+#endif // NK_TARGET_HASWELL
+
+#if NK_TARGET_SKYLAKE
+/** @copydoc nk_kld_f64 */
+NK_PUBLIC void nk_kld_f64_skylake(nk_f64_t const *a, nk_f64_t const *b, nk_size_t n, nk_f64_t *result);
+/** @copydoc nk_jsd_f64 */
+NK_PUBLIC void nk_jsd_f64_skylake(nk_f64_t const *a, nk_f64_t const *b, nk_size_t n, nk_f64_t *result);
+/** @copydoc nk_kld_f32 */
+NK_PUBLIC void nk_kld_f32_skylake(nk_f32_t const *a, nk_f32_t const *b, nk_size_t n, nk_f64_t *result);
+/** @copydoc nk_jsd_f32 */
+NK_PUBLIC void nk_jsd_f32_skylake(nk_f32_t const *a, nk_f32_t const *b, nk_size_t n, nk_f64_t *result);
+/** @copydoc nk_kld_f16 */
+NK_PUBLIC void nk_kld_f16_skylake(nk_f16_t const *a, nk_f16_t const *b, nk_size_t n, nk_f32_t *result);
+/** @copydoc nk_jsd_f16 */
+NK_PUBLIC void nk_jsd_f16_skylake(nk_f16_t const *a, nk_f16_t const *b, nk_size_t n, nk_f32_t *result);
+#endif // NK_TARGET_SKYLAKE
+
+#if NK_TARGET_RVV
+/** @copydoc nk_kld_f32 */
+NK_PUBLIC void nk_kld_f32_rvv(nk_f32_t const *a, nk_f32_t const *b, nk_size_t n, nk_f64_t *result);
+/** @copydoc nk_jsd_f32 */
+NK_PUBLIC void nk_jsd_f32_rvv(nk_f32_t const *a, nk_f32_t const *b, nk_size_t n, nk_f64_t *result);
+/** @copydoc nk_kld_f64 */
+NK_PUBLIC void nk_kld_f64_rvv(nk_f64_t const *a, nk_f64_t const *b, nk_size_t n, nk_f64_t *result);
+/** @copydoc nk_jsd_f64 */
+NK_PUBLIC void nk_jsd_f64_rvv(nk_f64_t const *a, nk_f64_t const *b, nk_size_t n, nk_f64_t *result);
+/** @copydoc nk_kld_f16 */
+NK_PUBLIC void nk_kld_f16_rvv(nk_f16_t const *a, nk_f16_t const *b, nk_size_t n, nk_f32_t *result);
+/** @copydoc nk_jsd_f16 */
+NK_PUBLIC void nk_jsd_f16_rvv(nk_f16_t const *a, nk_f16_t const *b, nk_size_t n, nk_f32_t *result);
+/** @copydoc nk_kld_bf16 */
+NK_PUBLIC void nk_kld_bf16_rvv(nk_bf16_t const *a, nk_bf16_t const *b, nk_size_t n, nk_f32_t *result);
+/** @copydoc nk_jsd_bf16 */
+NK_PUBLIC void nk_jsd_bf16_rvv(nk_bf16_t const *a, nk_bf16_t const *b, nk_size_t n, nk_f32_t *result);
+#endif // NK_TARGET_RVV
+
+/** @brief Returns the output dtype for probability measures (KLD, JSD). */
+NK_INTERNAL nk_dtype_t nk_probability_output_dtype(nk_dtype_t dtype) {
+    switch (dtype) {
+    case nk_f64_k: return nk_f64_k;
+    case nk_f32_k: return nk_f64_k;
+    case nk_f16_k: return nk_f32_k;
+    case nk_bf16_k: return nk_f32_k;
+    default: return nk_dtype_unknown_k;
+    }
+}
+
+#if defined(__cplusplus)
+} // extern "C"
+#endif
+
+#include "numkong/probability/serial.h"
+#include "numkong/probability/neon.h"
+#include "numkong/probability/haswell.h"
+#include "numkong/probability/skylake.h"
+#include "numkong/probability/rvv.h"
+
+#if defined(__cplusplus)
+extern "C" {
+#endif
+
+#if !NK_DYNAMIC_DISPATCH
+
+NK_PUBLIC void nk_kld_f16(nk_f16_t const *a, nk_f16_t const *b, nk_size_t n, nk_f32_t *result) {
+#if NK_TARGET_NEONHALF
+    nk_kld_f16_neonhalf(a, b, n, result);
+#elif NK_TARGET_SKYLAKE
+    nk_kld_f16_skylake(a, b, n, result);
+#elif NK_TARGET_HASWELL
+    nk_kld_f16_haswell(a, b, n, result);
+#elif NK_TARGET_RVV
+    nk_kld_f16_rvv(a, b, n, result);
+#else
+    nk_kld_f16_serial(a, b, n, result);
+#endif
+}
+
+NK_PUBLIC void nk_kld_bf16(nk_bf16_t const *a, nk_bf16_t const *b, nk_size_t n, nk_f32_t *result) {
+#if NK_TARGET_RVV
+    nk_kld_bf16_rvv(a, b, n, result);
+#else
+    nk_kld_bf16_serial(a, b, n, result);
+#endif
+}
+
+NK_PUBLIC void nk_kld_f32(nk_f32_t const *a, nk_f32_t const *b, nk_size_t n, nk_f64_t *result) {
+#if NK_TARGET_NEON
+    nk_kld_f32_neon(a, b, n, result);
+#elif NK_TARGET_SKYLAKE
+    nk_kld_f32_skylake(a, b, n, result);
+#elif NK_TARGET_RVV
+    nk_kld_f32_rvv(a, b, n, result);
+#else
+    nk_kld_f32_serial(a, b, n, result);
+#endif
+}
+
+NK_PUBLIC void nk_kld_f64(nk_f64_t const *a, nk_f64_t const *b, nk_size_t n, nk_f64_t *result) {
+#if NK_TARGET_SKYLAKE
+    nk_kld_f64_skylake(a, b, n, result);
+#elif NK_TARGET_HASWELL
+    nk_kld_f64_haswell(a, b, n, result);
+#elif NK_TARGET_RVV
+    nk_kld_f64_rvv(a, b, n, result);
+#else
+    nk_kld_f64_serial(a, b, n, result);
+#endif
+}
+
+NK_PUBLIC void nk_jsd_f16(nk_f16_t const *a, nk_f16_t const *b, nk_size_t n, nk_f32_t *result) {
+#if NK_TARGET_NEONHALF
+    nk_jsd_f16_neonhalf(a, b, n, result);
+#elif NK_TARGET_SKYLAKE
+    nk_jsd_f16_skylake(a, b, n, result);
+#elif NK_TARGET_HASWELL
+    nk_jsd_f16_haswell(a, b, n, result);
+#elif NK_TARGET_RVV
+    nk_jsd_f16_rvv(a, b, n, result);
+#else
+    nk_jsd_f16_serial(a, b, n, result);
+#endif
+}
+
+NK_PUBLIC void nk_jsd_bf16(nk_bf16_t const *a, nk_bf16_t const *b, nk_size_t n, nk_f32_t *result) {
+#if NK_TARGET_RVV
+    nk_jsd_bf16_rvv(a, b, n, result);
+#else
+    nk_jsd_bf16_serial(a, b, n, result);
+#endif
+}
+
+NK_PUBLIC void nk_jsd_f32(nk_f32_t const *a, nk_f32_t const *b, nk_size_t n, nk_f64_t *result) {
+#if NK_TARGET_NEON
+    nk_jsd_f32_neon(a, b, n, result);
+#elif NK_TARGET_SKYLAKE
+    nk_jsd_f32_skylake(a, b, n, result);
+#elif NK_TARGET_RVV
+    nk_jsd_f32_rvv(a, b, n, result);
+#else
+    nk_jsd_f32_serial(a, b, n, result);
+#endif
+}
+
+NK_PUBLIC void nk_jsd_f64(nk_f64_t const *a, nk_f64_t const *b, nk_size_t n, nk_f64_t *result) {
+#if NK_TARGET_SKYLAKE
+    nk_jsd_f64_skylake(a, b, n, result);
+#elif NK_TARGET_HASWELL
+    nk_jsd_f64_haswell(a, b, n, result);
+#elif NK_TARGET_RVV
+    nk_jsd_f64_rvv(a, b, n, result);
+#else
+    nk_jsd_f64_serial(a, b, n, result);
+#endif
+}
+
+#endif // !NK_DYNAMIC_DISPATCH
+
+#if defined(__cplusplus)
+} // extern "C"
+#endif
+
+#endif

package/include/numkong/probability.hpp

@@ -0,0 +1,120 @@
+/**
+ *  @brief C++ wrappers for SIMD-accelerated Similarity Measures for Probability Distributions.
+ *  @file include/numkong/probability.hpp
+ *  @author Ash Vardanian
+ *  @date February 5, 2026
+ */
+#ifndef NK_PROBABILITY_HPP
+#define NK_PROBABILITY_HPP
+
+#include <cstdint>
+#include <type_traits>
+
+#include "numkong/probability.h"
+
+#include "numkong/types.hpp"
+
+namespace ashvardanian::numkong {
+
+/**
+ *  @brief Kullback-Leibler divergence: Σ pᵢ × log(pᵢ / qᵢ)
+ *  @param[in] p,q First and second probability distributions
+ *  @param[in] d Number of dimensions in input vectors
+ *  @param[out] r Pointer to output divergence value
+ *
+ *  @tparam in_type_ Input distribution type (probability vectors)
+ *  @tparam result_type_ Result type, defaults to `in_type_::probability_result_t`
+ *  @tparam allow_simd_ Enable SIMD kernel dispatch when `prefer_simd_k`
+ */
+template <numeric_dtype in_type_, numeric_dtype result_type_ = typename in_type_::probability_result_t,
+          allow_simd_t allow_simd_ = prefer_simd_k>
+void kld(in_type_ const *p, in_type_ const *q, std::size_t d, result_type_ *r) noexcept {
+    constexpr bool simd = allow_simd_ == prefer_simd_k &&
+                          std::is_same_v<result_type_, typename in_type_::probability_result_t>;
+
+    if constexpr (std::is_same_v<in_type_, f64_t> && simd) nk_kld_f64(&p->raw_, &q->raw_, d, &r->raw_);
+    else if constexpr (std::is_same_v<in_type_, f32_t> && simd) nk_kld_f32(&p->raw_, &q->raw_, d, &r->raw_);
+    else if constexpr (std::is_same_v<in_type_, f16_t> && simd) nk_kld_f16(&p->raw_, &q->raw_, d, &r->raw_);
+    else if constexpr (std::is_same_v<in_type_, bf16_t> && simd) nk_kld_bf16(&p->raw_, &q->raw_, d, &r->raw_);
+    // Scalar fallback
+    else {
+        result_type_ sum {};
+        for (std::size_t i = 0; i < d; i++) {
+            result_type_ pi(p[i]), qi(q[i]);
+            if (pi > result_type_(0)) sum = sum + pi * (pi / qi).log();
+        }
+        *r = sum;
+    }
+}
+
+/**
+ *  @brief Jensen-Shannon distance: √(½ × (KL(p‖m) + KL(q‖m))), where m = (p + q) / 2
+ *  @param[in] p,q First and second probability distributions
+ *  @param[in] d Number of dimensions in input vectors
+ *  @param[out] r Pointer to output distance value
+ *
+ *  @tparam in_type_ Input distribution type (probability vectors)
+ *  @tparam result_type_ Result type, defaults to `in_type_::probability_result_t`
+ *  @tparam allow_simd_ Enable SIMD kernel dispatch when `prefer_simd_k`
+ */
+template <numeric_dtype in_type_, numeric_dtype result_type_ = typename in_type_::probability_result_t,
+          allow_simd_t allow_simd_ = prefer_simd_k>
+void jsd(in_type_ const *p, in_type_ const *q, std::size_t d, result_type_ *r) noexcept {
+    constexpr bool simd = allow_simd_ == prefer_simd_k &&
+                          std::is_same_v<result_type_, typename in_type_::probability_result_t>;
+
+    if constexpr (std::is_same_v<in_type_, f64_t> && simd) nk_jsd_f64(&p->raw_, &q->raw_, d, &r->raw_);
+    else if constexpr (std::is_same_v<in_type_, f32_t> && simd) nk_jsd_f32(&p->raw_, &q->raw_, d, &r->raw_);
+    else if constexpr (std::is_same_v<in_type_, f16_t> && simd) nk_jsd_f16(&p->raw_, &q->raw_, d, &r->raw_);
+    else if constexpr (std::is_same_v<in_type_, bf16_t> && simd) nk_jsd_bf16(&p->raw_, &q->raw_, d, &r->raw_);
+    // Scalar fallback
+    else {
+        result_type_ sum {};
+        result_type_ half(0.5);
+        for (std::size_t i = 0; i < d; i++) {
+            result_type_ pi(p[i]), qi(q[i]);
+            result_type_ mi = half * (pi + qi);
+            if (pi > result_type_(0)) sum = sum + pi * (pi / mi).log();
+            if (qi > result_type_(0)) sum = sum + qi * (qi / mi).log();
+        }
+        // JSD distance = sqrt(divergence / 2), clamped to non-negative
+        result_type_ divergence = half * sum;
+        *r = divergence > result_type_(0) ? divergence.sqrt() : result_type_(0);
+    }
+}
+
+} // namespace ashvardanian::numkong
+
+#include "numkong/tensor.hpp"
+
+namespace ashvardanian::numkong {
+
+template <numeric_dtype in_type_, numeric_dtype result_type_ = typename in_type_::probability_result_t,
+          allow_simd_t allow_simd_ = prefer_simd_k, std::size_t max_rank_a_, std::size_t max_rank_b_>
+void kld(tensor_view<in_type_, max_rank_a_> p, tensor_view<in_type_, max_rank_b_> q, std::size_t d,
+         result_type_ *r) noexcept {
+    kld<in_type_, result_type_, allow_simd_>(p.data(), q.data(), d, r);
+}
+
+template <numeric_dtype in_type_, numeric_dtype result_type_ = typename in_type_::probability_result_t,
+          allow_simd_t allow_simd_ = prefer_simd_k>
+void kld(vector_view<in_type_> p, vector_view<in_type_> q, std::size_t d, result_type_ *r) noexcept {
+    kld<in_type_, result_type_, allow_simd_>(p.data(), q.data(), d, r);
+}
+
+template <numeric_dtype in_type_, numeric_dtype result_type_ = typename in_type_::probability_result_t,
+          allow_simd_t allow_simd_ = prefer_simd_k, std::size_t max_rank_a_, std::size_t max_rank_b_>
+void jsd(tensor_view<in_type_, max_rank_a_> p, tensor_view<in_type_, max_rank_b_> q, std::size_t d,
+         result_type_ *r) noexcept {
+    jsd<in_type_, result_type_, allow_simd_>(p.data(), q.data(), d, r);
+}
+
+template <numeric_dtype in_type_, numeric_dtype result_type_ = typename in_type_::probability_result_t,
+          allow_simd_t allow_simd_ = prefer_simd_k>
+void jsd(vector_view<in_type_> p, vector_view<in_type_> q, std::size_t d, result_type_ *r) noexcept {
+    jsd<in_type_, result_type_, allow_simd_>(p.data(), q.data(), d, r);
+}
+
+} // namespace ashvardanian::numkong
+
+#endif // NK_PROBABILITY_HPP

package/include/numkong/random.h

@@ -0,0 +1,50 @@
+/**
+ *  @brief SIMD-accelerated Pseudo-Random Number Generators.
+ *  @file include/numkong/random.h
+ *  @author Ash Vardanian
+ *  @date January 11, 2026
+ *
+ *  Implements following statistical distributions
+ *
+ *  - Uniform Distribution
+ *  - Gaussian (Normal) Distribution
+ *
+ *  For dtypes:
+ *
+ *  - 64-bit floating point numbers
+ *  - 32-bit floating point numbers
+ *  - 16-bit floating point numbers
+ *  - 16-bit brain-floating point numbers
+ *  - 8-bit floating point numbers
+ *  - 8-bit integers
+ *
+ *  For hardware architectures:
+ *
+ *  - Arm: NEON, SSVE
+ *  - x86: Haswell, Ice Lake, Skylake, Genoa
+ *
+ *  @section usage Usage and Benefits
+ *
+ *
+ *
+ *  @section references References
+ *
+ *  - x86 intrinsics: https://www.intel.com/content/www/us/en/docs/intrinsics-guide/
+ *  - Arm intrinsics: https://developer.arm.com/architectures/instruction-sets/intrinsics/
+ *
+ */
+#ifndef NK_RANDOM_H
+#define NK_RANDOM_H
+
+#include "numkong/types.h"
+#include "numkong/cast.h"
+
+#if defined(__cplusplus)
+extern "C" {
+#endif // defined(__cplusplus)
+
+#if defined(__cplusplus)
+} // extern "C"
+#endif // defined(__cplusplus)
+
+#endif // NK_RANDOM_H