ipfrs_tensorlogic/
activation.rs

1//! Activation functions with forward and backward passes for tensor computations.
2//!
3//! Provides element-wise activation functions (ReLU, LeakyReLU, Sigmoid, Tanh,
4//! Softmax, GELU, Swish) with both forward evaluation and gradient (backward)
5//! computation.
6
7use std::f64::consts::PI;
8
9/// Type of activation function.
10#[derive(Debug, Clone, Copy, PartialEq, Eq)]
11pub enum ActivationType {
12    /// Rectified Linear Unit: max(0, x)
13    ReLU,
14    /// Leaky ReLU: x if x > 0, alpha * x otherwise
15    LeakyReLU,
16    /// Logistic sigmoid: 1 / (1 + exp(-x))
17    Sigmoid,
18    /// Hyperbolic tangent
19    Tanh,
20    /// Softmax (applied across all elements)
21    Softmax,
22    /// Gaussian Error Linear Unit (approximate form)
23    GELU,
24    /// Swish / SiLU: x * sigmoid(x)
25    Swish,
26}
27
28/// Configuration for an activation function.
29#[derive(Debug, Clone)]
30pub struct ActivationConfig {
31    /// Which activation function to use.
32    pub activation_type: ActivationType,
33    /// Alpha parameter for LeakyReLU (default 0.01).
34    pub leaky_alpha: f64,
35}
36
37impl Default for ActivationConfig {
38    fn default() -> Self {
39        Self {
40            activation_type: ActivationType::ReLU,
41            leaky_alpha: 0.01,
42        }
43    }
44}
45
46/// Runtime statistics for a `TensorActivation` instance.
47#[derive(Debug, Clone)]
48pub struct ActivationStats {
49    /// The activation type in use.
50    pub activation_type: ActivationType,
51    /// Number of forward passes executed.
52    pub forward_calls: u64,
53    /// Number of backward passes executed.
54    pub backward_calls: u64,
55}
56
57/// Activation layer with forward/backward support and call statistics.
58pub struct TensorActivation {
59    config: ActivationConfig,
60    forward_calls: u64,
61    backward_calls: u64,
62}
63
64impl TensorActivation {
65    /// Create a new activation layer from the given configuration.
66    pub fn new(config: ActivationConfig) -> Self {
67        Self {
68            config,
69            forward_calls: 0,
70            backward_calls: 0,
71        }
72    }
73
74    // ------------------------------------------------------------------
75    // Static helpers
76    // ------------------------------------------------------------------
77
78    /// ReLU activation: max(0, x).
79    #[inline]
80    pub fn relu(x: f64) -> f64 {
81        if x > 0.0 {
82            x
83        } else {
84            0.0
85        }
86    }
87
88    /// Sigmoid activation: 1 / (1 + exp(-x)).
89    #[inline]
90    pub fn sigmoid(x: f64) -> f64 {
91        if x >= 0.0 {
92            let e = (-x).exp();
93            1.0 / (1.0 + e)
94        } else {
95            let e = x.exp();
96            e / (1.0 + e)
97        }
98    }
99
100    /// Tanh activation (wrapper around `f64::tanh`).
101    #[inline]
102    pub fn tanh_act(x: f64) -> f64 {
103        x.tanh()
104    }
105
106    /// GELU activation (approximate):
107    /// 0.5 * x * (1 + tanh(sqrt(2/pi) * (x + 0.044715 * x^3)))
108    #[inline]
109    pub fn gelu(x: f64) -> f64 {
110        let sqrt_2_over_pi = (2.0 / PI).sqrt();
111        let inner = sqrt_2_over_pi * (x + 0.044715 * x * x * x);
112        0.5 * x * (1.0 + inner.tanh())
113    }
114
115    /// Swish activation: x * sigmoid(x).
116    #[inline]
117    pub fn swish(x: f64) -> f64 {
118        x * Self::sigmoid(x)
119    }
120
121    /// Softmax over a slice of values, using the log-sum-exp trick for
122    /// numerical stability.
123    pub fn softmax(input: &[f64]) -> Vec<f64> {
124        if input.is_empty() {
125            return Vec::new();
126        }
127        let max_val = input.iter().copied().fold(f64::NEG_INFINITY, f64::max);
128        let exps: Vec<f64> = input.iter().map(|&x| (x - max_val).exp()).collect();
129        let sum: f64 = exps.iter().sum();
130        if sum == 0.0 {
131            // Degenerate case – return uniform.
132            let n = input.len() as f64;
133            return vec![1.0 / n; input.len()];
134        }
135        exps.iter().map(|&e| e / sum).collect()
136    }
137
138    // ------------------------------------------------------------------
139    // Forward pass
140    // ------------------------------------------------------------------
141
142    /// Apply the configured activation element-wise (Softmax across all
143    /// elements).
144    pub fn forward(&mut self, input: &[f64]) -> Vec<f64> {
145        self.forward_calls += 1;
146        match self.config.activation_type {
147            ActivationType::ReLU => input.iter().map(|&x| Self::relu(x)).collect(),
148            ActivationType::LeakyReLU => {
149                let alpha = self.config.leaky_alpha;
150                input
151                    .iter()
152                    .map(|&x| if x > 0.0 { x } else { alpha * x })
153                    .collect()
154            }
155            ActivationType::Sigmoid => input.iter().map(|&x| Self::sigmoid(x)).collect(),
156            ActivationType::Tanh => input.iter().map(|&x| Self::tanh_act(x)).collect(),
157            ActivationType::Softmax => Self::softmax(input),
158            ActivationType::GELU => input.iter().map(|&x| Self::gelu(x)).collect(),
159            ActivationType::Swish => input.iter().map(|&x| Self::swish(x)).collect(),
160        }
161    }
162
163    // ------------------------------------------------------------------
164    // Backward pass
165    // ------------------------------------------------------------------
166
167    /// Compute the gradient of the loss w.r.t. the activation input.
168    ///
169    /// * `input`       – the original pre-activation values.
170    /// * `grad_output` – the upstream gradient (dL/d(output)).
171    ///
172    /// Returns dL/d(input).
173    pub fn backward(&mut self, input: &[f64], grad_output: &[f64]) -> Vec<f64> {
174        self.backward_calls += 1;
175
176        match self.config.activation_type {
177            ActivationType::ReLU => input
178                .iter()
179                .zip(grad_output.iter())
180                .map(|(&x, &g)| if x > 0.0 { g } else { 0.0 })
181                .collect(),
182            ActivationType::LeakyReLU => {
183                let alpha = self.config.leaky_alpha;
184                input
185                    .iter()
186                    .zip(grad_output.iter())
187                    .map(|(&x, &g)| if x > 0.0 { g } else { alpha * g })
188                    .collect()
189            }
190            ActivationType::Sigmoid => input
191                .iter()
192                .zip(grad_output.iter())
193                .map(|(&x, &g)| {
194                    let s = Self::sigmoid(x);
195                    g * s * (1.0 - s)
196                })
197                .collect(),
198            ActivationType::Tanh => input
199                .iter()
200                .zip(grad_output.iter())
201                .map(|(&x, &g)| {
202                    let t = x.tanh();
203                    g * (1.0 - t * t)
204                })
205                .collect(),
206            ActivationType::GELU => input
207                .iter()
208                .zip(grad_output.iter())
209                .map(|(&x, &g)| g * Self::gelu_derivative(x))
210                .collect(),
211            ActivationType::Swish => input
212                .iter()
213                .zip(grad_output.iter())
214                .map(|(&x, &g)| {
215                    let sw = Self::swish(x);
216                    let sig = Self::sigmoid(x);
217                    g * (sw + sig * (1.0 - sw))
218                })
219                .collect(),
220            ActivationType::Softmax => {
221                // Jacobian-vector product for softmax:
222                // dL/dx_i = s_i * (g_i - dot(g, s))
223                let s = Self::softmax(input);
224                let dot: f64 = grad_output
225                    .iter()
226                    .zip(s.iter())
227                    .map(|(&g, &si)| g * si)
228                    .sum();
229                s.iter()
230                    .zip(grad_output.iter())
231                    .map(|(&si, &gi)| si * (gi - dot))
232                    .collect()
233            }
234        }
235    }
236
237    // ------------------------------------------------------------------
238    // Stats
239    // ------------------------------------------------------------------
240
241    /// Return runtime statistics.
242    pub fn stats(&self) -> ActivationStats {
243        ActivationStats {
244            activation_type: self.config.activation_type,
245            forward_calls: self.forward_calls,
246            backward_calls: self.backward_calls,
247        }
248    }
249
250    // ------------------------------------------------------------------
251    // Private helpers
252    // ------------------------------------------------------------------
253
254    /// Approximate derivative of GELU.
255    ///
256    /// d/dx GELU(x) ≈ 0.5 * (1 + tanh(u)) + 0.5 * x * sech²(u) * u'
257    /// where u = sqrt(2/π) * (x + 0.044715 x³), u' = sqrt(2/π) * (1 + 3*0.044715 x²)
258    #[inline]
259    fn gelu_derivative(x: f64) -> f64 {
260        let sqrt_2_over_pi = (2.0 / PI).sqrt();
261        let x3 = x * x * x;
262        let u = sqrt_2_over_pi * (x + 0.044715 * x3);
263        let tanh_u = u.tanh();
264        let sech2_u = 1.0 - tanh_u * tanh_u;
265        let u_prime = sqrt_2_over_pi * (1.0 + 3.0 * 0.044715 * x * x);
266        0.5 * (1.0 + tanh_u) + 0.5 * x * sech2_u * u_prime
267    }
268}
269
270// ======================================================================
271// Tests
272// ======================================================================
273
274#[cfg(test)]
275mod tests {
276    use super::*;
277
278    // Helper to build a TensorActivation quickly.
279    fn make(ty: ActivationType) -> TensorActivation {
280        TensorActivation::new(ActivationConfig {
281            activation_type: ty,
282            leaky_alpha: 0.01,
283        })
284    }
285
286    // ------------------------------------------------------------------
287    // ReLU
288    // ------------------------------------------------------------------
289
290    #[test]
291    fn relu_forward_positive() {
292        let mut act = make(ActivationType::ReLU);
293        let out = act.forward(&[1.0, 2.0, 3.0]);
294        assert_eq!(out, vec![1.0, 2.0, 3.0]);
295    }
296
297    #[test]
298    fn relu_forward_zeros_negatives() {
299        let mut act = make(ActivationType::ReLU);
300        let out = act.forward(&[-1.0, -0.5, 0.0, 0.5]);
301        assert_eq!(out, vec![0.0, 0.0, 0.0, 0.5]);
302    }
303
304    #[test]
305    fn relu_backward() {
306        let mut act = make(ActivationType::ReLU);
307        let grad = act.backward(&[-1.0, 0.0, 1.0], &[1.0, 1.0, 1.0]);
308        assert_eq!(grad, vec![0.0, 0.0, 1.0]);
309    }
310
311    #[test]
312    fn relu_static_helper() {
313        assert_eq!(TensorActivation::relu(5.0), 5.0);
314        assert_eq!(TensorActivation::relu(-3.0), 0.0);
315        assert_eq!(TensorActivation::relu(0.0), 0.0);
316    }
317
318    // ------------------------------------------------------------------
319    // LeakyReLU
320    // ------------------------------------------------------------------
321
322    #[test]
323    fn leaky_relu_forward() {
324        let mut act = make(ActivationType::LeakyReLU);
325        let out = act.forward(&[-10.0, 0.0, 5.0]);
326        assert!((out[0] - (-0.1)).abs() < 1e-12);
327        assert_eq!(out[1], 0.0);
328        assert_eq!(out[2], 5.0);
329    }
330
331    #[test]
332    fn leaky_relu_backward() {
333        let mut act = make(ActivationType::LeakyReLU);
334        let grad = act.backward(&[-2.0, 3.0], &[1.0, 1.0]);
335        assert!((grad[0] - 0.01).abs() < 1e-12);
336        assert_eq!(grad[1], 1.0);
337    }
338
339    // ------------------------------------------------------------------
340    // Sigmoid
341    // ------------------------------------------------------------------
342
343    #[test]
344    fn sigmoid_forward_range() {
345        let mut act = make(ActivationType::Sigmoid);
346        let out = act.forward(&[-100.0, -1.0, 0.0, 1.0, 100.0]);
347        for &v in &out {
348            assert!((0.0..=1.0).contains(&v), "sigmoid out of range: {v}");
349        }
350        assert!((out[2] - 0.5).abs() < 1e-12, "sigmoid(0) should be 0.5");
351    }
352
353    #[test]
354    fn sigmoid_backward() {
355        let mut act = make(ActivationType::Sigmoid);
356        let grad = act.backward(&[0.0], &[1.0]);
357        // sigmoid(0) = 0.5, derivative = 0.5 * 0.5 = 0.25
358        assert!((grad[0] - 0.25).abs() < 1e-12);
359    }
360
361    #[test]
362    fn sigmoid_static_helper() {
363        assert!((TensorActivation::sigmoid(0.0) - 0.5).abs() < 1e-12);
364        assert!(TensorActivation::sigmoid(100.0) > 0.999);
365        assert!(TensorActivation::sigmoid(-100.0) < 0.001);
366    }
367
368    // ------------------------------------------------------------------
369    // Tanh
370    // ------------------------------------------------------------------
371
372    #[test]
373    fn tanh_forward_range() {
374        let mut act = make(ActivationType::Tanh);
375        let out = act.forward(&[-100.0, -1.0, 0.0, 1.0, 100.0]);
376        for &v in &out {
377            assert!((-1.0..=1.0).contains(&v), "tanh out of range: {v}");
378        }
379        assert!(out[2].abs() < 1e-12, "tanh(0) should be 0");
380    }
381
382    #[test]
383    fn tanh_backward() {
384        let mut act = make(ActivationType::Tanh);
385        let grad = act.backward(&[0.0], &[1.0]);
386        // tanh(0) = 0, derivative = 1 - 0^2 = 1
387        assert!((grad[0] - 1.0).abs() < 1e-12);
388    }
389
390    #[test]
391    fn tanh_static_helper() {
392        assert!(TensorActivation::tanh_act(0.0).abs() < 1e-12);
393        assert!((TensorActivation::tanh_act(1.0) - 1.0_f64.tanh()).abs() < 1e-14);
394    }
395
396    // ------------------------------------------------------------------
397    // Softmax
398    // ------------------------------------------------------------------
399
400    #[test]
401    fn softmax_sums_to_one() {
402        let mut act = make(ActivationType::Softmax);
403        let out = act.forward(&[1.0, 2.0, 3.0, 4.0]);
404        let sum: f64 = out.iter().sum();
405        assert!(
406            (sum - 1.0).abs() < 1e-12,
407            "softmax sum should be 1, got {sum}"
408        );
409    }
410
411    #[test]
412    fn softmax_monotonicity() {
413        let out = TensorActivation::softmax(&[1.0, 2.0, 3.0]);
414        assert!(out[0] < out[1] && out[1] < out[2]);
415    }
416
417    #[test]
418    fn softmax_backward() {
419        let mut act = make(ActivationType::Softmax);
420        let input = vec![1.0, 2.0, 3.0];
421        let grad_out = vec![1.0, 0.0, 0.0];
422        let grad = act.backward(&input, &grad_out);
423        // Sum of softmax backward gradients should be 0
424        let grad_sum: f64 = grad.iter().sum();
425        assert!(
426            grad_sum.abs() < 1e-12,
427            "softmax grad sum should be ~0, got {grad_sum}"
428        );
429    }
430
431    #[test]
432    fn softmax_static_helper() {
433        let out = TensorActivation::softmax(&[0.0, 0.0, 0.0]);
434        for &v in &out {
435            assert!((v - 1.0 / 3.0).abs() < 1e-12);
436        }
437    }
438
439    // ------------------------------------------------------------------
440    // GELU
441    // ------------------------------------------------------------------
442
443    #[test]
444    fn gelu_forward_approximation() {
445        let mut act = make(ActivationType::GELU);
446        let out = act.forward(&[0.0, 1.0, -1.0]);
447        // GELU(0) = 0
448        assert!(out[0].abs() < 1e-12);
449        // GELU(1) ≈ 0.8412
450        assert!((out[1] - 0.8412).abs() < 0.001);
451        // GELU(-1) ≈ -0.1588
452        assert!((out[2] - (-0.1588)).abs() < 0.001);
453    }
454
455    #[test]
456    fn gelu_backward() {
457        let mut act = make(ActivationType::GELU);
458        let grad = act.backward(&[0.0], &[1.0]);
459        // GELU'(0) = 0.5
460        assert!((grad[0] - 0.5).abs() < 1e-6);
461    }
462
463    #[test]
464    fn gelu_static_helper() {
465        assert!(TensorActivation::gelu(0.0).abs() < 1e-12);
466    }
467
468    // ------------------------------------------------------------------
469    // Swish
470    // ------------------------------------------------------------------
471
472    #[test]
473    fn swish_forward() {
474        let mut act = make(ActivationType::Swish);
475        let out = act.forward(&[0.0, 1.0, -1.0]);
476        // swish(0) = 0 * 0.5 = 0
477        assert!(out[0].abs() < 1e-12);
478        // swish(1) = 1 * sigmoid(1)
479        let expected = TensorActivation::sigmoid(1.0);
480        assert!((out[1] - expected).abs() < 1e-12);
481    }
482
483    #[test]
484    fn swish_backward() {
485        let mut act = make(ActivationType::Swish);
486        let grad = act.backward(&[0.0], &[1.0]);
487        // swish(0)=0, sigmoid(0)=0.5 => derivative = 0 + 0.5*(1-0) = 0.5
488        assert!((grad[0] - 0.5).abs() < 1e-12);
489    }
490
491    #[test]
492    fn swish_static_helper() {
493        assert!(TensorActivation::swish(0.0).abs() < 1e-12);
494        let s5 = TensorActivation::swish(5.0);
495        assert!((s5 - 5.0 * TensorActivation::sigmoid(5.0)).abs() < 1e-12);
496    }
497
498    // ------------------------------------------------------------------
499    // Empty input
500    // ------------------------------------------------------------------
501
502    #[test]
503    fn empty_input_forward() {
504        let mut act = make(ActivationType::ReLU);
505        assert!(act.forward(&[]).is_empty());
506    }
507
508    #[test]
509    fn empty_input_backward() {
510        let mut act = make(ActivationType::Sigmoid);
511        assert!(act.backward(&[], &[]).is_empty());
512    }
513
514    #[test]
515    fn empty_softmax() {
516        assert!(TensorActivation::softmax(&[]).is_empty());
517    }
518
519    // ------------------------------------------------------------------
520    // Stats tracking
521    // ------------------------------------------------------------------
522
523    #[test]
524    fn stats_tracking() {
525        let mut act = make(ActivationType::GELU);
526        let s = act.stats();
527        assert_eq!(s.forward_calls, 0);
528        assert_eq!(s.backward_calls, 0);
529        assert_eq!(s.activation_type, ActivationType::GELU);
530
531        act.forward(&[1.0]);
532        act.forward(&[2.0]);
533        act.backward(&[1.0], &[1.0]);
534
535        let s = act.stats();
536        assert_eq!(s.forward_calls, 2);
537        assert_eq!(s.backward_calls, 1);
538    }
539
540    // ------------------------------------------------------------------
541    // Numerical gradient checks
542    // ------------------------------------------------------------------
543
544    fn numerical_gradient(act: &mut TensorActivation, x: f64, eps: f64) -> f64 {
545        let mut act_plus = make(act.stats().activation_type);
546        let mut act_minus = make(act.stats().activation_type);
547        let f_plus = act_plus.forward(&[x + eps])[0];
548        let f_minus = act_minus.forward(&[x - eps])[0];
549        (f_plus - f_minus) / (2.0 * eps)
550    }
551
552    #[test]
553    fn numerical_grad_relu() {
554        let mut act = make(ActivationType::ReLU);
555        let analytic = act.backward(&[1.0], &[1.0])[0];
556        let numeric = numerical_gradient(&mut act, 1.0, 1e-5);
557        assert!((analytic - numeric).abs() < 1e-4);
558    }
559
560    #[test]
561    fn numerical_grad_sigmoid() {
562        let mut act = make(ActivationType::Sigmoid);
563        for &x in &[-2.0, -0.5, 0.0, 0.5, 2.0] {
564            let analytic = act.backward(&[x], &[1.0])[0];
565            let numeric = numerical_gradient(&mut act, x, 1e-5);
566            assert!(
567                (analytic - numeric).abs() < 1e-4,
568                "sigmoid grad mismatch at x={x}: analytic={analytic}, numeric={numeric}"
569            );
570        }
571    }
572
573    #[test]
574    fn numerical_grad_tanh() {
575        let mut act = make(ActivationType::Tanh);
576        for &x in &[-2.0, 0.0, 1.5] {
577            let analytic = act.backward(&[x], &[1.0])[0];
578            let numeric = numerical_gradient(&mut act, x, 1e-5);
579            assert!(
580                (analytic - numeric).abs() < 1e-4,
581                "tanh grad mismatch at x={x}"
582            );
583        }
584    }
585
586    #[test]
587    fn numerical_grad_gelu() {
588        let mut act = make(ActivationType::GELU);
589        for &x in &[-2.0, -0.5, 0.0, 0.5, 2.0] {
590            let analytic = act.backward(&[x], &[1.0])[0];
591            let numeric = numerical_gradient(&mut act, x, 1e-5);
592            assert!(
593                (analytic - numeric).abs() < 1e-3,
594                "gelu grad mismatch at x={x}: analytic={analytic}, numeric={numeric}"
595            );
596        }
597    }
598
599    #[test]
600    fn numerical_grad_swish() {
601        let mut act = make(ActivationType::Swish);
602        for &x in &[-2.0, -0.5, 0.0, 0.5, 2.0] {
603            let analytic = act.backward(&[x], &[1.0])[0];
604            let numeric = numerical_gradient(&mut act, x, 1e-5);
605            assert!(
606                (analytic - numeric).abs() < 1e-4,
607                "swish grad mismatch at x={x}: analytic={analytic}, numeric={numeric}"
608            );
609        }
610    }
611
612    #[test]
613    fn numerical_grad_leaky_relu() {
614        let mut act = make(ActivationType::LeakyReLU);
615        for &x in &[-2.0, 2.0] {
616            let analytic = act.backward(&[x], &[1.0])[0];
617            let numeric = numerical_gradient(&mut act, x, 1e-5);
618            assert!(
619                (analytic - numeric).abs() < 1e-4,
620                "leaky_relu grad mismatch at x={x}"
621            );
622        }
623    }
624
625    // ------------------------------------------------------------------
626    // Default config
627    // ------------------------------------------------------------------
628
629    #[test]
630    fn default_config() {
631        let cfg = ActivationConfig::default();
632        assert_eq!(cfg.activation_type, ActivationType::ReLU);
633        assert!((cfg.leaky_alpha - 0.01).abs() < 1e-14);
634    }
635}
ipfrs_tensorlogic/activation.rs

ipfrs_tensorlogic/
activation.rs
