scirs2_neural/
activations_minimal.rs

1//! Minimal activation functions without Layer trait dependencies
2//!
3//! This module provides SIMD-accelerated implementations for f32/f64 1D arrays,
4//! with automatic fallback to generic implementations for other types/dimensions.
5
6use crate::error::Result;
7use scirs2_core::ndarray::{Array, ArrayView1, IxDyn, Zip};
8use scirs2_core::numeric::Float;
9use scirs2_core::simd_ops::SimdUnifiedOps;
10use std::fmt::Debug;
11
12/// Trait for activation functions
13pub trait Activation<F> {
14    /// Forward pass of the activation function
15    fn forward(
16        &self,
17        input: &Array<F, scirs2_core::ndarray::IxDyn>,
18    ) -> Result<Array<F, scirs2_core::ndarray::IxDyn>>;
19
20    /// Backward pass of the activation function
21    fn backward(
22        &self,
23        grad_output: &Array<F, scirs2_core::ndarray::IxDyn>,
24        input: &Array<F, scirs2_core::ndarray::IxDyn>,
25    ) -> Result<Array<F, scirs2_core::ndarray::IxDyn>>;
26}
27
28/// GELU activation function
29#[derive(Debug, Clone, Copy)]
30pub struct GELU {
31    fast: bool,
32}
33
34impl GELU {
35    pub fn new() -> Self {
36        Self { fast: false }
37    }
38
39    pub fn fast() -> Self {
40        Self { fast: true }
41    }
42}
43
44impl Default for GELU {
45    fn default() -> Self {
46        Self::new()
47    }
48}
49
50impl<F: Float + Debug> Activation<F> for GELU {
51    fn forward(
52        &self,
53        input: &Array<F, scirs2_core::ndarray::IxDyn>,
54    ) -> Result<Array<F, scirs2_core::ndarray::IxDyn>> {
55        let mut output = input.clone();
56
57        if self.fast {
58            let sqrt_2_over_pi =
59                F::from(0.7978845608028654).expect("Failed to convert constant to float");
60            let coeff = F::from(0.044715).expect("Failed to convert constant to float");
61            let half = F::from(0.5).expect("Failed to convert constant to float");
62            let one = F::one();
63
64            Zip::from(&mut output).for_each(|x| {
65                let x3 = *x * *x * *x;
66                let inner = sqrt_2_over_pi * (*x + coeff * x3);
67                *x = half * *x * (one + inner.tanh());
68            });
69        } else {
70            let sqrt_pi_over_2 =
71                F::from(1.2533141373155).expect("Failed to convert constant to float");
72            let coeff = F::from(0.044715).expect("Failed to convert constant to float");
73            let half = F::from(0.5).expect("Failed to convert constant to float");
74            let one = F::one();
75
76            Zip::from(&mut output).for_each(|x| {
77                let x2 = *x * *x;
78                let inner = sqrt_pi_over_2 * *x * (one + coeff * x2);
79                *x = half * *x * (one + inner.tanh());
80            });
81        }
82
83        Ok(output)
84    }
85
86    fn backward(
87        &self,
88        grad_output: &Array<F, scirs2_core::ndarray::IxDyn>,
89        input: &Array<F, scirs2_core::ndarray::IxDyn>,
90    ) -> Result<Array<F, scirs2_core::ndarray::IxDyn>> {
91        let mut grad_input = Array::zeros(grad_output.raw_dim());
92
93        if self.fast {
94            let sqrt_2_over_pi =
95                F::from(0.7978845608028654).expect("Failed to convert constant to float");
96            let coeff = F::from(0.044715).expect("Failed to convert constant to float");
97            let half = F::from(0.5).expect("Failed to convert constant to float");
98            let one = F::one();
99            let three = F::from(3.0).expect("Failed to convert constant to float");
100
101            Zip::from(&mut grad_input)
102                .and(grad_output)
103                .and(input)
104                .for_each(|grad_in, &grad_out, &x| {
105                    let x2 = x * x;
106                    let x3 = x2 * x;
107                    let inner = sqrt_2_over_pi * (x + coeff * x3);
108                    let tanh_inner = inner.tanh();
109                    let sech_sq = one - tanh_inner * tanh_inner;
110                    let d_inner_dx = sqrt_2_over_pi * (one + three * coeff * x2);
111                    let dgelu_dx = half * (one + tanh_inner) + half * x * sech_sq * d_inner_dx;
112                    *grad_in = grad_out * dgelu_dx;
113                });
114        } else {
115            let sqrt_pi_over_2 =
116                F::from(1.2533141373155).expect("Failed to convert constant to float");
117            let coeff = F::from(0.044715).expect("Failed to convert constant to float");
118            let half = F::from(0.5).expect("Failed to convert constant to float");
119            let one = F::one();
120            let three = F::from(3.0).expect("Failed to convert constant to float");
121
122            Zip::from(&mut grad_input)
123                .and(grad_output)
124                .and(input)
125                .for_each(|grad_in, &grad_out, &x| {
126                    let x2 = x * x;
127                    let inner = sqrt_pi_over_2 * x * (one + coeff * x2);
128                    let tanh_inner = inner.tanh();
129                    let sech_sq = one - tanh_inner * tanh_inner;
130                    let d_inner_dx = sqrt_pi_over_2 * (one + three * coeff * x2);
131                    let dgelu_dx = half * (one + tanh_inner) + half * x * sech_sq * d_inner_dx;
132                    *grad_in = grad_out * dgelu_dx;
133                });
134        }
135
136        Ok(grad_input)
137    }
138}
139
140/// Tanh activation function
141#[derive(Debug, Clone, Copy)]
142pub struct Tanh;
143
144impl Tanh {
145    pub fn new() -> Self {
146        Self
147    }
148}
149
150impl Default for Tanh {
151    fn default() -> Self {
152        Self::new()
153    }
154}
155
156impl<F: Float + Debug> Activation<F> for Tanh {
157    fn forward(
158        &self,
159        input: &Array<F, scirs2_core::ndarray::IxDyn>,
160    ) -> Result<Array<F, scirs2_core::ndarray::IxDyn>> {
161        let mut output = input.clone();
162        Zip::from(&mut output).for_each(|x| {
163            *x = x.tanh();
164        });
165        Ok(output)
166    }
167
168    fn backward(
169        &self,
170        grad_output: &Array<F, scirs2_core::ndarray::IxDyn>,
171        input: &Array<F, scirs2_core::ndarray::IxDyn>,
172    ) -> Result<Array<F, scirs2_core::ndarray::IxDyn>> {
173        let mut grad_input = Array::zeros(grad_output.raw_dim());
174
175        Zip::from(&mut grad_input)
176            .and(grad_output)
177            .and(input)
178            .for_each(|grad_in, &grad_out, &x| {
179                let tanh_x = x.tanh();
180                let derivative = F::one() - tanh_x * tanh_x;
181                *grad_in = grad_out * derivative;
182            });
183
184        Ok(grad_input)
185    }
186}
187
188/// Sigmoid activation function
189#[derive(Debug, Clone, Copy)]
190pub struct Sigmoid;
191
192impl Sigmoid {
193    pub fn new() -> Self {
194        Self
195    }
196}
197
198impl Default for Sigmoid {
199    fn default() -> Self {
200        Self::new()
201    }
202}
203
204impl<F: Float + Debug> Activation<F> for Sigmoid {
205    fn forward(
206        &self,
207        input: &Array<F, scirs2_core::ndarray::IxDyn>,
208    ) -> Result<Array<F, scirs2_core::ndarray::IxDyn>> {
209        let mut output = input.clone();
210        let one = F::one();
211        Zip::from(&mut output).for_each(|x| {
212            *x = one / (one + (-*x).exp());
213        });
214        Ok(output)
215    }
216
217    fn backward(
218        &self,
219        grad_output: &Array<F, scirs2_core::ndarray::IxDyn>,
220        input: &Array<F, scirs2_core::ndarray::IxDyn>,
221    ) -> Result<Array<F, scirs2_core::ndarray::IxDyn>> {
222        let mut grad_input = Array::zeros(grad_output.raw_dim());
223        let one = F::one();
224
225        Zip::from(&mut grad_input)
226            .and(grad_output)
227            .and(input)
228            .for_each(|grad_in, &grad_out, &x| {
229                let sigmoid_x = one / (one + (-x).exp());
230                let derivative = sigmoid_x * (one - sigmoid_x);
231                *grad_in = grad_out * derivative;
232            });
233
234        Ok(grad_input)
235    }
236}
237
238/// ReLU activation function
239#[derive(Debug, Clone, Copy)]
240pub struct ReLU {
241    alpha: f64,
242}
243
244impl ReLU {
245    pub fn new() -> Self {
246        Self { alpha: 0.0 }
247    }
248
249    pub fn leaky(alpha: f64) -> Self {
250        Self { alpha }
251    }
252
253    /// Try SIMD path for f64 arrays
254    #[inline]
255    fn try_simd_f64<F: Float>(&self, input: &Array<F, IxDyn>, alpha: F) -> Option<Array<F, IxDyn>> {
256        // Check if F is f64 using size_of
257        if std::mem::size_of::<F>() != std::mem::size_of::<f64>() {
258            return None;
259        }
260
261        // SAFETY: We've verified the type size matches f64
262        unsafe {
263            let ptr = input.as_ptr() as *const f64;
264            let len = input.len();
265            let slice = std::slice::from_raw_parts(ptr, len);
266            let arr1d = ArrayView1::from(slice);
267
268            let alpha_f64 = *(&alpha as *const F as *const f64);
269            let result = if alpha_f64 == 0.0 {
270                f64::simd_relu(&arr1d)
271            } else {
272                f64::simd_leaky_relu(&arr1d, alpha_f64)
273            };
274
275            // Convert back to F
276            let result_ptr = result.as_ptr() as *const F;
277            let result_slice = std::slice::from_raw_parts(result_ptr, result.len());
278            let result_dyn =
279                Array::from_shape_vec(IxDyn(&[result.len()]), result_slice.to_vec()).ok()?;
280
281            Some(result_dyn)
282        }
283    }
284
285    /// Try SIMD path for f32 arrays
286    #[inline]
287    fn try_simd_f32<F: Float>(&self, input: &Array<F, IxDyn>, alpha: F) -> Option<Array<F, IxDyn>> {
288        // Check if F is f32 using size_of
289        if std::mem::size_of::<F>() != std::mem::size_of::<f32>() {
290            return None;
291        }
292
293        // SAFETY: We've verified the type size matches f32
294        unsafe {
295            let ptr = input.as_ptr() as *const f32;
296            let len = input.len();
297            let slice = std::slice::from_raw_parts(ptr, len);
298            let arr1d = ArrayView1::from(slice);
299
300            let alpha_f32 = *(&alpha as *const F as *const f32);
301            let result = if alpha_f32 == 0.0 {
302                f32::simd_relu(&arr1d)
303            } else {
304                f32::simd_leaky_relu(&arr1d, alpha_f32)
305            };
306
307            // Convert back to F
308            let result_ptr = result.as_ptr() as *const F;
309            let result_slice = std::slice::from_raw_parts(result_ptr, result.len());
310            let result_dyn =
311                Array::from_shape_vec(IxDyn(&[result.len()]), result_slice.to_vec()).ok()?;
312
313            Some(result_dyn)
314        }
315    }
316}
317
318impl Default for ReLU {
319    fn default() -> Self {
320        Self::new()
321    }
322}
323
324impl<F: Float + Debug> Activation<F> for ReLU {
325    fn forward(
326        &self,
327        input: &Array<F, scirs2_core::ndarray::IxDyn>,
328    ) -> Result<Array<F, scirs2_core::ndarray::IxDyn>> {
329        let zero = F::zero();
330        let alpha = F::from(self.alpha).unwrap_or(zero);
331
332        // SIMD fast path for 1D f32/f64 arrays
333        if input.ndim() == 1 {
334            // Try f64 SIMD path
335            if let Some(result) = self.try_simd_f64(input, alpha) {
336                return Ok(result);
337            }
338            // Try f32 SIMD path
339            if let Some(result) = self.try_simd_f32(input, alpha) {
340                return Ok(result);
341            }
342        }
343
344        // Generic fallback for all other cases
345        let mut output = input.clone();
346        Zip::from(&mut output).for_each(|x| {
347            if *x < zero {
348                *x = alpha * *x;
349            }
350        });
351        Ok(output)
352    }
353
354    fn backward(
355        &self,
356        grad_output: &Array<F, scirs2_core::ndarray::IxDyn>,
357        input: &Array<F, scirs2_core::ndarray::IxDyn>,
358    ) -> Result<Array<F, scirs2_core::ndarray::IxDyn>> {
359        let mut grad_input = Array::zeros(grad_output.raw_dim());
360        let zero = F::zero();
361        let one = F::one();
362        let alpha = F::from(self.alpha).unwrap_or(zero);
363
364        Zip::from(&mut grad_input)
365            .and(grad_output)
366            .and(input)
367            .for_each(|grad_in, &grad_out, &x| {
368                let derivative = if x > zero { one } else { alpha };
369                *grad_in = grad_out * derivative;
370            });
371
372        Ok(grad_input)
373    }
374}
375
376/// Softmax activation function
377#[derive(Debug, Clone, Copy)]
378pub struct Softmax {
379    axis: isize,
380}
381
382impl Softmax {
383    pub fn new(axis: isize) -> Self {
384        Self { axis }
385    }
386
387    /// Try SIMD path for f64 arrays
388    #[inline]
389    fn try_simd_f64<F: Float>(&self, input: &Array<F, IxDyn>) -> Option<Array<F, IxDyn>> {
390        // Check if F is f64 using size_of
391        if std::mem::size_of::<F>() != std::mem::size_of::<f64>() {
392            return None;
393        }
394
395        // SAFETY: We've verified the type size matches f64
396        unsafe {
397            let ptr = input.as_ptr() as *const f64;
398            let len = input.len();
399            let slice = std::slice::from_raw_parts(ptr, len);
400            let arr1d = ArrayView1::from(slice);
401
402            let result = f64::simd_softmax(&arr1d);
403
404            // Convert back to F
405            let result_ptr = result.as_ptr() as *const F;
406            let result_slice = std::slice::from_raw_parts(result_ptr, result.len());
407            let result_dyn =
408                Array::from_shape_vec(IxDyn(&[result.len()]), result_slice.to_vec()).ok()?;
409
410            Some(result_dyn)
411        }
412    }
413
414    /// Try SIMD path for f32 arrays
415    #[inline]
416    fn try_simd_f32<F: Float>(&self, input: &Array<F, IxDyn>) -> Option<Array<F, IxDyn>> {
417        // Check if F is f32 using size_of
418        if std::mem::size_of::<F>() != std::mem::size_of::<f32>() {
419            return None;
420        }
421
422        // SAFETY: We've verified the type size matches f32
423        unsafe {
424            let ptr = input.as_ptr() as *const f32;
425            let len = input.len();
426            let slice = std::slice::from_raw_parts(ptr, len);
427            let arr1d = ArrayView1::from(slice);
428
429            let result = f32::simd_softmax(&arr1d);
430
431            // Convert back to F
432            let result_ptr = result.as_ptr() as *const F;
433            let result_slice = std::slice::from_raw_parts(result_ptr, result.len());
434            let result_dyn =
435                Array::from_shape_vec(IxDyn(&[result.len()]), result_slice.to_vec()).ok()?;
436
437            Some(result_dyn)
438        }
439    }
440}
441
442impl Default for Softmax {
443    fn default() -> Self {
444        Self::new(-1)
445    }
446}
447
448impl<F: Float + Debug> Activation<F> for Softmax {
449    fn forward(
450        &self,
451        input: &Array<F, scirs2_core::ndarray::IxDyn>,
452    ) -> Result<Array<F, scirs2_core::ndarray::IxDyn>> {
453        // SIMD fast path for 1D f32/f64 arrays
454        if input.ndim() == 1 {
455            // Try f64 SIMD path
456            if let Some(result) = self.try_simd_f64(input) {
457                return Ok(result);
458            }
459            // Try f32 SIMD path
460            if let Some(result) = self.try_simd_f32(input) {
461                return Ok(result);
462            }
463        }
464
465        // Generic fallback for multi-dimensional or other types
466        let mut output = input.clone();
467
468        // Simple softmax implementation for the last axis
469        if self.axis == -1 || self.axis as usize == input.ndim() - 1 {
470            // For 1D case or applying to last axis
471            let max_val = input.fold(F::neg_infinity(), |acc, &x| if x > acc { x } else { acc });
472
473            // Subtract max for numerical stability
474            Zip::from(&mut output).for_each(|x| {
475                *x = (*x - max_val).exp();
476            });
477
478            // Sum all exponentials
479            let sum = output.sum();
480
481            // Normalize
482            Zip::from(&mut output).for_each(|x| {
483                *x = *x / sum;
484            });
485        }
486
487        Ok(output)
488    }
489
490    fn backward(
491        &self,
492        grad_output: &Array<F, scirs2_core::ndarray::IxDyn>,
493        input: &Array<F, scirs2_core::ndarray::IxDyn>,
494    ) -> Result<Array<F, scirs2_core::ndarray::IxDyn>> {
495        // Forward pass to get softmax _output
496        let softmax_output = self.forward(input)?;
497        let mut grad_input = Array::zeros(grad_output.raw_dim());
498
499        // For softmax: grad = softmax * (grad_out - (softmax * grad_out).sum())
500        let sum_grad = Zip::from(&softmax_output)
501            .and(grad_output)
502            .fold(F::zero(), |acc, &s, &g| acc + s * g);
503
504        Zip::from(&mut grad_input)
505            .and(&softmax_output)
506            .and(grad_output)
507            .for_each(|grad_in, &s, &grad_out| {
508                *grad_in = s * (grad_out - sum_grad);
509            });
510
511        Ok(grad_input)
512    }
513}
514#[cfg(test)]
515mod tests {
516    use super::*;
517    use scirs2_core::ndarray::Array;
518
519    #[test]
520    fn test_relu_simd_f64_1d() {
521        let relu = ReLU::new();
522        let input = Array::from_vec(vec![-2.0, -1.0, 0.0, 1.0, 2.0]).into_dyn();
523        let output = relu.forward(&input).expect("Operation failed");
524        let expected = Array::from_vec(vec![0.0, 0.0, 0.0, 1.0, 2.0]).into_dyn();
525        assert_eq!(output, expected);
526    }
527
528    #[test]
529    fn test_relu_simd_f32_1d() {
530        let relu = ReLU::new();
531        let input = Array::from_vec(vec![-2.0f32, -1.0, 0.0, 1.0, 2.0]).into_dyn();
532        let output = relu.forward(&input).expect("Operation failed");
533        let expected = Array::from_vec(vec![0.0f32, 0.0, 0.0, 1.0, 2.0]).into_dyn();
534        assert_eq!(output, expected);
535    }
536
537    #[test]
538    fn test_leaky_relu_simd_f64_1d() {
539        let relu = ReLU::leaky(0.01);
540        let input = Array::from_vec(vec![-2.0, -1.0, 0.0, 1.0, 2.0]).into_dyn();
541        let output = relu.forward(&input).expect("Operation failed");
542        let expected = Array::from_vec(vec![-0.02, -0.01, 0.0, 1.0, 2.0]).into_dyn();
543        for (a, b) in output.iter().zip(expected.iter()) {
544            assert!((a - b).abs() < 1e-10);
545        }
546    }
547
548    #[test]
549    fn test_leaky_relu_simd_f32_1d() {
550        let relu = ReLU::leaky(0.01);
551        let input = Array::from_vec(vec![-2.0f32, -1.0, 0.0, 1.0, 2.0]).into_dyn();
552        let output = relu.forward(&input).expect("Operation failed");
553        let expected = Array::from_vec(vec![-0.02f32, -0.01, 0.0, 1.0, 2.0]).into_dyn();
554        for (a, b) in output.iter().zip(expected.iter()) {
555            assert!((a - b).abs() < 1e-6);
556        }
557    }
558
559    #[test]
560    fn test_relu_large_array_f64() {
561        // Test with large array to verify SIMD path is used
562        let relu = ReLU::new();
563        let input: Vec<f64> = (0..10000).map(|i| i as f64 - 5000.0).collect();
564        let input_arr = Array::from_vec(input.clone()).into_dyn();
565        let output = relu.forward(&input_arr).expect("Operation failed");
566
567        // Verify correctness
568        for (i, &val) in output.iter().enumerate() {
569            let expected = if input[i] > 0.0 { input[i] } else { 0.0 };
570            assert!((val - expected).abs() < 1e-10);
571        }
572    }
573
574    #[test]
575    fn test_leaky_relu_large_array_f32() {
576        // Test with large array to verify SIMD path is used
577        let relu = ReLU::leaky(0.01);
578        let input: Vec<f32> = (0..10000).map(|i| i as f32 - 5000.0).collect();
579        let input_arr = Array::from_vec(input.clone()).into_dyn();
580        let output = relu.forward(&input_arr).expect("Operation failed");
581
582        // Verify correctness
583        for (i, &val) in output.iter().enumerate() {
584            let expected = if input[i] > 0.0 {
585                input[i]
586            } else {
587                0.01 * input[i]
588            };
589            assert!((val - expected).abs() < 1e-5);
590        }
591    }
592
593    #[test]
594    fn test_relu_2d_fallback() {
595        // Test that 2D arrays still work (using generic fallback)
596        let relu = ReLU::new();
597        let input = Array::from_shape_vec((2, 3), vec![-2.0, -1.0, 0.0, 1.0, 2.0, 3.0])
598            .expect("Operation failed")
599            .into_dyn();
600        let output = relu.forward(&input).expect("Operation failed");
601        let expected = Array::from_shape_vec((2, 3), vec![0.0, 0.0, 0.0, 1.0, 2.0, 3.0])
602            .expect("Operation failed")
603            .into_dyn();
604        assert_eq!(output, expected);
605    }
606
607    #[test]
608    fn test_relu_backward() {
609        let relu = ReLU::new();
610        let input = Array::from_vec(vec![-1.0, 0.0, 1.0]).into_dyn();
611        let grad_output = Array::from_vec(vec![1.0, 1.0, 1.0]).into_dyn();
612        let grad_input = relu
613            .backward(&grad_output, &input)
614            .expect("Operation failed");
615        let expected = Array::from_vec(vec![0.0, 0.0, 1.0]).into_dyn();
616        assert_eq!(grad_input, expected);
617    }
618
619    #[test]
620    fn test_leaky_relu_backward() {
621        let relu = ReLU::leaky(0.01);
622        let input = Array::from_vec(vec![-1.0, 0.0, 1.0]).into_dyn();
623        let grad_output = Array::from_vec(vec![1.0, 1.0, 1.0]).into_dyn();
624        let grad_input = relu
625            .backward(&grad_output, &input)
626            .expect("Operation failed");
627        let expected = Array::from_vec(vec![0.01, 0.01, 1.0]).into_dyn();
628        for (a, b) in grad_input.iter().zip(expected.iter()) {
629            assert!((a - b).abs() < 1e-10);
630        }
631    }
632}
scirs2_neural/activations_minimal.rs

scirs2_neural/
activations_minimal.rs
