singe_kernel/cpu/
normalization.rs

1//! RMS norm, layer norm, group norm, sparsemax, and related reductions.
2
3pub fn sparsemax(input: &[f32], rows: usize, cols: usize) -> Vec<f32> {
4    let mut out = vec![0.0f32; rows * cols];
5    for row in 0..rows {
6        let values = &input[row * cols..(row + 1) * cols];
7        let mut sorted = values.to_vec();
8        sorted.sort_by(|lhs, rhs| rhs.partial_cmp(lhs).unwrap());
9        let mut prefix_sum = 0.0f32;
10        let mut support = 0usize;
11        for (index, &value) in sorted.iter().enumerate() {
12            prefix_sum += value;
13            let rank = index + 1;
14            if 1.0 + rank as f32 * value > prefix_sum {
15                support = rank;
16            }
17        }
18        let tau = (sorted[..support].iter().sum::<f32>() - 1.0) / support as f32;
19        for col in 0..cols {
20            out[row * cols + col] = (values[col] - tau).max(0.0);
21        }
22    }
23    out
24}
25
26pub fn group_norm(
27    input: &[f32],
28    weight: &[f32],
29    bias: &[f32],
30    batch: usize,
31    channels: usize,
32    groups: usize,
33    spatial_len: usize,
34    eps: f32,
35) -> Vec<f32> {
36    let mut out = vec![0.0f32; batch * channels * spatial_len];
37    let channels_per_group = channels / groups;
38    for batch_index in 0..batch {
39        for group in 0..groups {
40            let mut sum = 0.0f32;
41            let mut sum_sq = 0.0f32;
42            for channel_in_group in 0..channels_per_group {
43                let channel = group * channels_per_group + channel_in_group;
44                for spatial in 0..spatial_len {
45                    let index = (batch_index * channels + channel) * spatial_len + spatial;
46                    let value = input[index];
47                    sum += value;
48                    sum_sq += value * value;
49                }
50            }
51            let group_len = (channels_per_group * spatial_len) as f32;
52            let mean = sum / group_len;
53            let variance = sum_sq / group_len - mean * mean;
54            let rstd = 1.0 / (variance + eps).sqrt();
55            for channel_in_group in 0..channels_per_group {
56                let channel = group * channels_per_group + channel_in_group;
57                for spatial in 0..spatial_len {
58                    let index = (batch_index * channels + channel) * spatial_len + spatial;
59                    out[index] = (input[index] - mean) * rstd * weight[channel] + bias[channel];
60                }
61            }
62        }
63    }
64    out
65}
singe_kernel/cpu/normalization.rs

singe_kernel/cpu/
normalization.rs
