boostr 0.2.0

ML framework built on numr - attention, quantization, model architectures
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132

//! Snake activation.
//!
//! Per Liu et al. 2020 ("Neural Networks Fail to Learn Periodic Functions and
//! How to Fix It"):
//!
//! ```text
//!     snake(x) = x + (1 / (α + ε)) * sin²(α · x)
//! ```
//!
//! Kokoro's ISTFTNet generator uses a learnable per-channel α shaped
//! `[1, C, 1]` in every `AdaINResBlock1` site (stored as `alpha1` / `alpha2`
//! `ParameterList`s in the state_dict). `ε` is a tiny positive constant that
//! keeps the reciprocal bounded when α is near zero — we follow the
//! `SnakeBeta` upstream in using `1e-9`.
//!
//! Composite op: no new kernel needed. Uses mul / sin / div / add from numr.

use crate::error::{Error, Result};
use numr::dtype::DType;
use numr::ops::{BinaryOps, ScalarOps, UnaryOps};
use numr::runtime::{Runtime, RuntimeClient};
use numr::tensor::Tensor;

/// Apply snake activation with a per-channel α.
///
/// * `x` — activation of shape `[B, C, T]`.
/// * `alpha` — shape `[1, C, 1]` (matches the state_dict layout used by
///   `torch.nn.ParameterList([Parameter(torch.ones(1, ch, 1))])`).
/// * `eps` — denominator floor; pass `1e-9` to match upstream, or a larger
///   value if numerical stability is a concern.
pub fn snake<R, C>(client: &C, x: &Tensor<R>, alpha: &Tensor<R>, eps: f64) -> Result<Tensor<R>>
where
    R: Runtime<DType = DType>,
    C: RuntimeClient<R> + BinaryOps<R> + UnaryOps<R> + ScalarOps<R>,
{
    let x_shape = x.shape();
    if x_shape.len() != 3 {
        return Err(Error::InvalidArgument {
            arg: "x",
            reason: format!("snake expects rank-3 [B, C, T], got {x_shape:?}"),
        });
    }
    let channels = x_shape[1];
    let alpha_shape = alpha.shape();
    if alpha_shape != [1, channels, 1] {
        return Err(Error::InvalidArgument {
            arg: "alpha",
            reason: format!("snake alpha shape must be [1, {channels}, 1], got {alpha_shape:?}"),
        });
    }

    // α · x (broadcast across B and T).
    let ax = client.mul(alpha, x).map_err(Error::Numr)?;
    // sin²(α · x)
    let s = client.sin(&ax).map_err(Error::Numr)?;
    let s_sq = client.mul(&s, &s).map_err(Error::Numr)?;
    // 1 / (α + ε) — add scalar ε then reciprocate.
    let alpha_eps = client.add_scalar(alpha, eps).map_err(Error::Numr)?;
    let inv = client.recip(&alpha_eps).map_err(Error::Numr)?;
    // x + (1 / (α + ε)) · sin²(α · x)
    let scaled = client.mul(&inv, &s_sq).map_err(Error::Numr)?;
    client.add(x, &scaled).map_err(Error::Numr)
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::test_utils::cpu_setup;
    use numr::runtime::cpu::CpuRuntime;

    #[test]
    fn zero_input_stays_zero() {
        let (client, device) = cpu_setup();
        let x = Tensor::<CpuRuntime>::from_slice(&[0.0f32; 6], &[1, 2, 3], &device);
        let alpha = Tensor::<CpuRuntime>::from_slice(&[1.0f32, 1.0], &[1, 2, 1], &device);
        let y = snake(&client, &x, &alpha, 1e-9).unwrap();
        for v in y.to_vec::<f32>() {
            assert!(v.abs() < 1e-6, "got {v}");
        }
    }

    #[test]
    fn matches_formula_for_alpha_1() {
        // snake(x) = x + sin²(x). Check pointwise.
        let (client, device) = cpu_setup();
        let input: Vec<f32> = vec![-1.5, -0.5, 0.0, 0.5, 1.5];
        let x = Tensor::<CpuRuntime>::from_slice(&input, &[1, 1, 5], &device);
        let alpha = Tensor::<CpuRuntime>::from_slice(&[1.0f32], &[1, 1, 1], &device);
        let y: Vec<f32> = snake(&client, &x, &alpha, 0.0).unwrap().to_vec();
        for (got, x_in) in y.iter().zip(&input) {
            let expected = x_in + x_in.sin().powi(2);
            assert!(
                (got - expected).abs() < 1e-5,
                "snake({x_in}) = {got} vs {expected}"
            );
        }
    }

    #[test]
    fn alpha_scales_frequency() {
        // For α=2, period halves: snake(x) = x + sin²(2x).
        let (client, device) = cpu_setup();
        let input: Vec<f32> = vec![0.5, 1.0];
        let x = Tensor::<CpuRuntime>::from_slice(&input, &[1, 1, 2], &device);
        let alpha = Tensor::<CpuRuntime>::from_slice(&[2.0f32], &[1, 1, 1], &device);
        let y: Vec<f32> = snake(&client, &x, &alpha, 0.0).unwrap().to_vec();
        for (got, x_in) in y.iter().zip(&input) {
            let expected = x_in + 0.5 * (2.0 * x_in).sin().powi(2);
            assert!(
                (got - expected).abs() < 1e-5,
                "snake α=2({x_in}) = {got} vs {expected}"
            );
        }
    }

    #[test]
    fn rejects_bad_alpha_shape() {
        let (client, device) = cpu_setup();
        let x = Tensor::<CpuRuntime>::from_slice(&[0.0f32; 6], &[1, 2, 3], &device);
        // Should be [1, 2, 1]; give [2] instead.
        let alpha = Tensor::<CpuRuntime>::from_slice(&[1.0f32, 1.0], &[2], &device);
        assert!(snake(&client, &x, &alpha, 1e-9).is_err());
    }

    #[test]
    fn rejects_bad_input_rank() {
        let (client, device) = cpu_setup();
        let x = Tensor::<CpuRuntime>::from_slice(&[0.0f32; 4], &[2, 2], &device);
        let alpha = Tensor::<CpuRuntime>::from_slice(&[1.0f32, 1.0], &[1, 2, 1], &device);
        assert!(snake(&client, &x, &alpha, 1e-9).is_err());
    }
}