Skip to content

Commit

Permalink
Adding burn::nn::Sigmoid (#2031)
Browse files Browse the repository at this point in the history
  • Loading branch information
RuelYasa authored Jul 17, 2024
1 parent ed8a91d commit 9804bf8
Show file tree
Hide file tree
Showing 2 changed files with 40 additions and 0 deletions.
2 changes: 2 additions & 0 deletions crates/burn-core/src/nn/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -29,6 +29,7 @@ mod prelu;
mod relu;
mod rnn;
mod rope_encoding;
mod sigmoid;
mod swiglu;
mod tanh;
mod unfold;
Expand All @@ -46,6 +47,7 @@ pub use prelu::*;
pub use relu::*;
pub use rnn::*;
pub use rope_encoding::*;
pub use sigmoid::*;
pub use swiglu::*;
pub use tanh::*;
pub use unfold::*;
38 changes: 38 additions & 0 deletions crates/burn-core/src/nn/sigmoid.rs
Original file line number Diff line number Diff line change
@@ -0,0 +1,38 @@
use crate as burn;

use crate::module::Module;
use crate::tensor::backend::Backend;
use crate::tensor::Tensor;

/// Applies the sigmoid function element-wise
/// See also [sigmoid](burn::tensor::activation::sigmoid)
#[derive(Module, Clone, Debug, Default)]
pub struct Sigmoid;

impl Sigmoid {
/// Create the module.
pub fn new() -> Self {
Self {}
}
/// Applies the forward pass on the input tensor.
///
/// # Shapes
///
/// - input: `[..., any]`
/// - output: `[..., any]`
pub fn forward<B: Backend, const D: usize>(&self, input: Tensor<B, D>) -> Tensor<B, D> {
crate::tensor::activation::sigmoid(input)
}
}

#[cfg(test)]
mod tests {
use super::*;

#[test]
fn display() {
let layer = Sigmoid::new();

assert_eq!(alloc::format!("{}", layer), "Sigmoid");
}
}

0 comments on commit 9804bf8

Please sign in to comment.