Skip to main content

hodu_nn/modules/
activation.rs

1use crate::{
2    compat::*,
3    module::Module,
4    state::{get_state, State},
5};
6use hodu_core::{error::HoduResult, scalar::Scalar, tensor::Tensor};
7
8#[derive(Module, Clone, Default)]
9pub struct ReLU;
10
11impl ReLU {
12    pub fn new() -> Self {
13        Self
14    }
15
16    pub fn forward(&self, input: &Tensor) -> HoduResult<Tensor> {
17        input.relu()
18    }
19
20    pub fn parameters(&mut self) -> Vec<&mut Tensor> {
21        vec![]
22    }
23}
24
25#[derive(Module, Clone, Default)]
26pub struct Sigmoid;
27
28impl Sigmoid {
29    pub fn new() -> Self {
30        Self
31    }
32
33    pub fn forward(&self, input: &Tensor) -> HoduResult<Tensor> {
34        input.sigmoid()
35    }
36
37    pub fn parameters(&mut self) -> Vec<&mut Tensor> {
38        vec![]
39    }
40}
41
42#[derive(Module, Clone, Default)]
43pub struct Tanh;
44
45impl Tanh {
46    pub fn new() -> Self {
47        Self
48    }
49
50    pub fn forward(&self, input: &Tensor) -> HoduResult<Tensor> {
51        input.tanh()
52    }
53
54    pub fn parameters(&mut self) -> Vec<&mut Tensor> {
55        vec![]
56    }
57}
58
59#[derive(Module, Clone, Default)]
60pub struct Gelu;
61
62impl Gelu {
63    pub fn new() -> Self {
64        Self
65    }
66
67    pub fn forward(&self, input: &Tensor) -> HoduResult<Tensor> {
68        input.gelu()
69    }
70
71    pub fn parameters(&mut self) -> Vec<&mut Tensor> {
72        vec![]
73    }
74}
75
76#[derive(Module, Clone, Default)]
77pub struct Softplus;
78
79impl Softplus {
80    pub fn new() -> Self {
81        Self
82    }
83
84    pub fn forward(&self, input: &Tensor) -> HoduResult<Tensor> {
85        input.softplus()
86    }
87
88    pub fn parameters(&mut self) -> Vec<&mut Tensor> {
89        vec![]
90    }
91}
92
93#[derive(Module, Clone, Default)]
94pub struct SiLU;
95
96impl SiLU {
97    pub fn new() -> Self {
98        Self
99    }
100
101    pub fn forward(&self, input: &Tensor) -> HoduResult<Tensor> {
102        input.silu()
103    }
104
105    pub fn parameters(&mut self) -> Vec<&mut Tensor> {
106        vec![]
107    }
108}
109
110#[derive(Module, Clone, Default)]
111pub struct Swish;
112
113impl Swish {
114    pub fn new() -> Self {
115        Self
116    }
117
118    pub fn forward(&self, input: &Tensor) -> HoduResult<Tensor> {
119        input.swish()
120    }
121
122    pub fn parameters(&mut self) -> Vec<&mut Tensor> {
123        vec![]
124    }
125}
126
127#[derive(Module, Clone, Default)]
128pub struct Mish;
129
130impl Mish {
131    pub fn new() -> Self {
132        Self
133    }
134
135    pub fn forward(&self, input: &Tensor) -> HoduResult<Tensor> {
136        input.mish()
137    }
138
139    pub fn parameters(&mut self) -> Vec<&mut Tensor> {
140        vec![]
141    }
142}
143
144#[derive(Module, Clone)]
145pub struct LeakyReLU {
146    exponent: Scalar,
147}
148
149impl LeakyReLU {
150    pub fn new(exponent: impl Into<Scalar>) -> Self {
151        Self {
152            exponent: exponent.into(),
153        }
154    }
155
156    pub fn forward(&self, input: &Tensor) -> HoduResult<Tensor> {
157        let exponent = self.exponent.to_dtype(input.dtype());
158        input.leaky_relu(exponent)
159    }
160
161    pub fn parameters(&mut self) -> Vec<&mut Tensor> {
162        vec![]
163    }
164}
165
166#[derive(Module, Clone)]
167pub struct ELU {
168    exponent: Scalar,
169}
170
171impl ELU {
172    pub fn new(exponent: impl Into<Scalar>) -> Self {
173        Self {
174            exponent: exponent.into(),
175        }
176    }
177
178    pub fn forward(&self, input: &Tensor) -> HoduResult<Tensor> {
179        let exponent = self.exponent.to_dtype(input.dtype());
180        input.elu(exponent)
181    }
182
183    pub fn parameters(&mut self) -> Vec<&mut Tensor> {
184        vec![]
185    }
186}
187
188#[derive(Module, Clone)]
189pub struct PReLU {
190    weight: Scalar,
191}
192
193impl PReLU {
194    pub fn new(weight: impl Into<Scalar>) -> Self {
195        Self { weight: weight.into() }
196    }
197
198    pub fn forward(&self, input: &Tensor) -> HoduResult<Tensor> {
199        let weight = self.weight.to_dtype(input.dtype());
200        input.prelu(weight)
201    }
202
203    pub fn parameters(&mut self) -> Vec<&mut Tensor> {
204        vec![]
205    }
206}
207
208#[derive(Module, Clone)]
209pub struct RReLU {
210    lower: Scalar,
211    upper: Scalar,
212}
213
214impl RReLU {
215    pub fn new(lower: impl Into<Scalar>, upper: impl Into<Scalar>) -> Self {
216        Self {
217            lower: lower.into(),
218            upper: upper.into(),
219        }
220    }
221
222    pub fn forward(&self, input: &Tensor) -> HoduResult<Tensor> {
223        let dtype = input.dtype();
224        let zero = Scalar::zero(dtype);
225
226        // Compute alpha based on training/evaluation mode
227        let alpha = if get_state() == State::Training {
228            // Training mode: use random alpha for each element
229            let lower_f32 = self.lower.to_f32();
230            let upper_f32 = self.upper.to_f32();
231            Tensor::rand_uniform_like(input, lower_f32, upper_f32)?
232        } else {
233            // Evaluation mode: use average of lower and upper bounds
234            let avg = (self.lower.to_f32() + self.upper.to_f32()) / 2.0;
235            let avg_scalar = Scalar::from_f32(avg, dtype);
236            Tensor::full_like(input, avg_scalar)?
237        };
238
239        // RReLU: x if x > 0, else alpha * x
240        let mask_pos = input.gt_scalar(zero)?;
241        let mask_neg = input.le_scalar(zero)?;
242        let positive_part = input.mul(&mask_pos.to_dtype(dtype)?)?;
243        let negative_part = input.mul(&alpha)?.mul(&mask_neg.to_dtype(dtype)?)?;
244
245        positive_part.add(&negative_part)
246    }
247
248    pub fn parameters(&mut self) -> Vec<&mut Tensor> {
249        vec![]
250    }
251}