1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
use burn::config::Config;
use burn::module::Module;
use burn::module::{Content, DisplaySettings, ModuleDisplay};
use burn::tensor::Tensor;
use burn_core as burn;
use burn::tensor::activation::threshold;
/// Threshold layer: returns `x` where `x > threshold`, and `value` otherwise.
///
/// Should be created with [ThresholdConfig](ThresholdConfig).
#[derive(Module, Debug)]
#[module(custom_display)]
pub struct Threshold {
/// The value to threshold at.
pub threshold: f64,
/// The value to replace with where the input is at or below the threshold.
pub value: f64,
}
/// Configuration to create a [Threshold](Threshold) layer using the [init function](ThresholdConfig::init).
#[derive(Config, Debug)]
pub struct ThresholdConfig {
/// The value to threshold at.
pub threshold: f64,
/// The value to replace with where the input is at or below the threshold.
pub value: f64,
}
impl ThresholdConfig {
/// Initialize a new [Threshold](Threshold) layer.
pub fn init(&self) -> Threshold {
Threshold {
threshold: self.threshold,
value: self.value,
}
}
}
impl ModuleDisplay for Threshold {
fn custom_settings(&self) -> Option<DisplaySettings> {
DisplaySettings::new()
.with_new_line_after_attribute(false)
.optional()
}
fn custom_content(&self, content: Content) -> Option<Content> {
content
.add("threshold", &self.threshold)
.add("value", &self.value)
.optional()
}
}
impl Threshold {
/// Forward pass for the Threshold layer.
///
/// See [threshold](burn::tensor::activation::threshold) for more information.
///
/// # Shapes
/// - input: `[..., any]`
/// - output: `[..., any]`
pub fn forward<const D: usize>(&self, input: Tensor<D>) -> Tensor<D> {
threshold(input, self.threshold, self.value)
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn display() {
let config = ThresholdConfig::new(1.0, 20.0).init();
assert_eq!(
alloc::format!("{config}"),
"Threshold {threshold: 1, value: 20}"
);
}
}