Tanh nn wrapper (tracel-ai#1903)

crates/burn-core/src/nn/mod.rs

@@ -30,6 +30,7 @@ mod relu;
 mod rnn;
 mod rope_encoding;
 mod swiglu;
+mod tanh;
 mod unfold;
 
 pub use dropout::*;
@@ -46,4 +47,5 @@ pub use relu::*;
 pub use rnn::*;
 pub use rope_encoding::*;
 pub use swiglu::*;
+pub use tanh::*;
 pub use unfold::*;

crates/burn-core/src/nn/tanh.rs

@@ -0,0 +1,26 @@
+use crate as burn;
+
+use crate::module::Module;
+use crate::tensor::backend::Backend;
+use crate::tensor::Tensor;
+
+/// Applies the tanh activation function element-wise
+/// See also [tanh](burn::tensor::activation::tanh)
+#[derive(Module, Clone, Debug, Default)]
+pub struct Tanh {}
+
+impl Tanh {
+ /// 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::tanh(input)
+ }
+}