needle_core/
ops.rs

1//! Primitive tensor operations: f32 only, no allocations.
2use crate::math;
3
4/// In-place softmax over a slice (stable via max subtraction).
5pub fn softmax_inplace(x: &mut [f32]) {
6    if x.is_empty() {
7        return;
8    }
9    let max = x.iter().cloned().fold(f32::NEG_INFINITY, f32::max);
10    // All-−∞ input: exp(−∞ − (−∞)) = exp(NaN) = NaN. Fall back to uniform.
11    if max == f32::NEG_INFINITY {
12        let uniform = 1.0 / x.len() as f32;
13        for v in x.iter_mut() {
14            *v = uniform;
15        }
16        return;
17    }
18    let mut sum = 0.0f32;
19    for v in x.iter_mut() {
20        *v = math::exp(*v - max);
21        sum += *v;
22    }
23    // sum > 0 is guaranteed (at least one exp(0) = 1.0 from the max element),
24    // but guard against pathological subnormal underflow.
25    if sum == 0.0 {
26        let uniform = 1.0 / x.len() as f32;
27        for v in x.iter_mut() {
28            *v = uniform;
29        }
30        return;
31    }
32    let inv = 1.0 / sum;
33    for v in x.iter_mut() {
34        *v *= inv;
35    }
36}
37
38/// Element-wise sigmoid: σ(x) = 1 / (1 + e^-x).
39#[inline]
40pub fn sigmoid(x: f32) -> f32 {
41    1.0 / (1.0 + math::exp(-x))
42}
43
44/// GELU activation (approximate, matches JAX nn.gelu default).
45/// Formula: 0.5 * x * (1 + tanh(sqrt(2/π) * (x + 0.044715 * x³)))
46#[inline]
47pub fn gelu(x: f32) -> f32 {
48    const C: f32 = 0.797_884_56; // sqrt(2/π)
49    const A: f32 = 0.044_715;
50    let inner = C * (x + A * x * x * x);
51    0.5 * x * (1.0 + math::tanh(inner))
52}
53
54/// Dot product of two equal-length slices.
55#[inline]
56pub fn dot(a: &[f32], b: &[f32]) -> f32 {
57    debug_assert_eq!(a.len(), b.len());
58    let mut acc = 0.0f32;
59    for (x, y) in a.iter().zip(b.iter()) {
60        acc += x * y;
61    }
62    acc
63}
64
65/// y += alpha * x  (SAXPY)
66#[inline]
67pub fn saxpy(alpha: f32, x: &[f32], y: &mut [f32]) {
68    debug_assert_eq!(x.len(), y.len());
69    for (yi, xi) in y.iter_mut().zip(x.iter()) {
70        *yi += alpha * xi;
71    }
72}
73
74/// Element-wise multiply: y[i] *= x[i]
75#[inline]
76pub fn elementwise_mul(x: &[f32], y: &mut [f32]) {
77    debug_assert_eq!(x.len(), y.len());
78    for (yi, xi) in y.iter_mut().zip(x.iter()) {
79        *yi *= xi;
80    }
81}
82
83/// Dense matmul: y[b,j] = Σ_i x[b,i] * w[i,j]
84/// w shape [in_dim, out_dim], x shape [batch, in_dim], y shape [batch, out_dim].
85pub fn matmul(x: &[f32], w: &[f32], y: &mut [f32], batch: usize, in_dim: usize, out_dim: usize) {
86    debug_assert_eq!(x.len(), batch * in_dim);
87    debug_assert_eq!(w.len(), in_dim * out_dim);
88    debug_assert_eq!(y.len(), batch * out_dim);
89    for b in 0..batch {
90        let xb = &x[b * in_dim..(b + 1) * in_dim];
91        let yb = &mut y[b * out_dim..(b + 1) * out_dim];
92        for j in 0..out_dim {
93            let mut acc = 0.0f32;
94            for i in 0..in_dim {
95                acc += xb[i] * w[i * out_dim + j];
96            }
97            yb[j] = acc;
98        }
99    }
100}
101
102/// Add bias in-place: x[b, j] += bias[j]
103pub fn add_bias(x: &mut [f32], bias: &[f32], batch: usize, out_dim: usize) {
104    debug_assert_eq!(x.len(), batch * out_dim);
105    debug_assert_eq!(bias.len(), out_dim);
106    for b in 0..batch {
107        for j in 0..out_dim {
108            x[b * out_dim + j] += bias[j];
109        }
110    }
111}
112
113/// SiLU / swish: `x * sigmoid(x)`.
114///
115/// Shared: v2 and v3 use it with identical semantics and it bakes in no
116/// geometry.
117#[inline]
118pub fn silu(x: f32) -> f32 {
119    x * sigmoid(x)
120}
needle_core/ops.rs

needle_core/
ops.rs
