burn_nn/activation/
soft_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::soft_shrink;
8use burn::tensor::backend::Backend;
9
10/// Soft Shrink layer.
11///
12/// Applies the Soft Shrink function element-wise:
13/// `soft_shrink(x) = x - lambda if x > lambda, x + lambda if x < -lambda, 0 otherwise`
14///
15/// Should be created with [SoftShrinkConfig](SoftShrinkConfig).
16#[derive(Module, Clone, Debug)]
17#[module(custom_display)]
18pub struct SoftShrink {
19    /// The lambda value for the Soft Shrink formulation.
20    pub lambda: f64,
21}
22
23/// Configuration to create a [SoftShrink](SoftShrink) layer using the [init function](SoftShrinkConfig::init).
24#[derive(Config, Debug)]
25pub struct SoftShrinkConfig {
26    /// The lambda value for the Soft Shrink formulation. Default is 0.5
27    #[config(default = "0.5")]
28    pub lambda: f64,
29}
30
31impl SoftShrinkConfig {
32    /// Initialize a new [SoftShrink](SoftShrink) Layer
33    pub fn init(&self) -> SoftShrink {
34        SoftShrink {
35            lambda: self.lambda,
36        }
37    }
38}
39
40impl ModuleDisplay for SoftShrink {
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 SoftShrink {
53    /// Forward pass for the Soft Shrink layer.
54    ///
55    /// See [soft_shrink](burn::tensor::activation::soft_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        soft_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_soft_shrink_forward() {
73        let device = Default::default();
74        let model: SoftShrink = SoftShrinkConfig::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, -0.5], [7.5, 0.0, 0.0]]);
79        assert_eq!(out.into_data(), expected);
80    }
81
82    #[test]
83    fn test_soft_shrink_with_lambda() {
84        let device = Default::default();
85        let model: SoftShrink = SoftShrinkConfig::new().with_lambda(0.25).init();
86        let input =
87            Tensor::<TestBackend, 2>::from_data([[0.125, -0.125, -0.5], [0.75, 0.1, 0.0]], &device);
88        let out = model.forward(input);
89        let expected = TensorData::from([[0.0_f32, 0.0, -0.25], [0.5, 0.0, 0.0]]);
90        assert_eq!(out.into_data(), expected);
91    }
92
93    #[test]
94    fn display() {
95        let config = SoftShrinkConfig::new().init();
96        assert_eq!(alloc::format!("{config}"), "SoftShrink {lambda: 0.5}");
97    }
98}
burn_nn/activation/soft_shrink.rs

burn_nn/activation/
soft_shrink.rs
