burn_nn/activation/
hard_shrink.rs

1use burn_core as burn;
2
3use burn::config::Config;
4use burn::module::Module;
5use burn::module::{Content, DisplaySettings, ModuleDisplay};
6use burn::tensor::Tensor;
7use burn::tensor::activation::hard_shrink;
8use burn::tensor::backend::Backend;
9
10/// Hard Shrink layer.
11///
12/// Applies the Hard Shrink function element-wise:
13/// `hard_shrink(x) = x if |x| > lambda else 0`
14///
15/// Should be created with [HardShrinkConfig](HardShrinkConfig).
16#[derive(Module, Clone, Debug)]
17#[module(custom_display)]
18pub struct HardShrink {
19    /// The lambda value for the Hard Shrink formulation.
20    pub lambda: f64,
21}
22
23/// Configuration to create a [HardShrink](HardShrink) layer using the [init function](HardShrinkConfig::init).
24#[derive(Config, Debug)]
25pub struct HardShrinkConfig {
26    /// The lambda value for the Hard Shrink formulation. Default is 0.5
27    #[config(default = "0.5")]
28    pub lambda: f64,
29}
30
31impl HardShrinkConfig {
32    /// Initialize a new [HardShrink](HardShrink) Layer
33    pub fn init(&self) -> HardShrink {
34        HardShrink {
35            lambda: self.lambda,
36        }
37    }
38}
39
40impl ModuleDisplay for HardShrink {
41    fn custom_settings(&self) -> Option<DisplaySettings> {
42        DisplaySettings::new()
43            .with_new_line_after_attribute(false)
44            .optional()
45    }
46
47    fn custom_content(&self, content: Content) -> Option<Content> {
48        content.add("lambda", &self.lambda).optional()
49    }
50}
51
52impl HardShrink {
53    /// Forward pass for the Hard Shrink layer.
54    ///
55    /// See [hard_shrink](burn::tensor::activation::hard_shrink) for more information.
56    ///
57    /// # Shapes
58    /// - input: `[..., any]`
59    /// - output: `[..., any]`
60    pub fn forward<B: Backend, const D: usize>(&self, input: Tensor<B, D>) -> Tensor<B, D> {
61        hard_shrink(input, self.lambda)
62    }
63}
64
65#[cfg(test)]
66mod tests {
67    use super::*;
68    use crate::TestBackend;
69    use burn::tensor::TensorData;
70
71    #[test]
72    fn test_hard_shrink_forward() {
73        let device = Default::default();
74        let model: HardShrink = HardShrinkConfig::new().init();
75        let input =
76            Tensor::<TestBackend, 2>::from_data([[0.5, -0.5, -1.0], [8.0, 0.3, 0.0]], &device);
77        let out = model.forward(input);
78        let expected = TensorData::from([[0.0_f32, 0.0, -1.0], [8.0, 0.0, 0.0]]);
79        assert_eq!(out.into_data(), expected);
80    }
81
82    #[test]
83    fn test_hard_shrink_with_lambda() {
84        let device = Default::default();
85        let model: HardShrink = HardShrinkConfig::new().with_lambda(0.2).init();
86        let input =
87            Tensor::<TestBackend, 2>::from_data([[0.1, -0.1, -0.3], [0.5, 0.1, 0.0]], &device);
88        let out = model.forward(input);
89        let expected = TensorData::from([[0.0_f32, 0.0, -0.3], [0.5, 0.0, 0.0]]);
90        assert_eq!(out.into_data(), expected);
91    }
92
93    #[test]
94    fn display() {
95        let config = HardShrinkConfig::new().init();
96        assert_eq!(alloc::format!("{config}"), "HardShrink {lambda: 0.5}");
97    }
98}
burn_nn/activation/hard_shrink.rs

burn_nn/activation/
hard_shrink.rs
