Repository: coreylowman/dfdx (language: rust)

Issue:
Add `tensor_ops::leaky_relu` and `nn::LeakyReLU`
See pytorch docs https://pytorch.org/docs/stable/generated/torch.nn.LeakyReLU.html#torch.nn.LeakyReLU.

Will need to figure out what the derivative of this function is, but that's probably online somewhere.

See https://github.com/coreylowman/dfdx/pull/397 for all the pieces you'll need to add for this.


Discussion / hints:
Hi, I'm going to take a stab at this!

Fix patch:
diff --git a/src/nn/activations.rs b/src/nn/activations.rs
index a3e0f7e52..b5ef6dc37 100644
--- a/src/nn/activations.rs
+++ b/src/nn/activations.rs
@@ -1,4 +1,9 @@
-use crate::{shapes::*, tensor::*, tensor_ops::*};
+use crate::{
+    prelude::{BuildModule, BuildOnDevice, TensorCollection, TensorOptions},
+    shapes::*,
+    tensor::*,
+    tensor_ops::*,
+};
 
 use super::module::{Module, NonMutableModule, ZeroSizedModule};
 
@@ -55,9 +60,159 @@ where
     }
 }
 
+/// Calls [prelu()] with constant value.
+#[derive(Debug, Clone, Copy)]
+pub struct LeakyReLU<E: Dtype>(pub E);
+
+impl<E: Dtype> Default for LeakyReLU<E> {
+    fn default() -> Self {
+        Self(E::from_f32(0.05).unwrap())
+    }
+}
+
+impl<E: Dtype> ZeroSizedModule for LeakyReLU<E> {}
+impl<E: Dtype> NonMutableModule for LeakyReLU<E> {}
+
+impl<S: ConstShape, E: Dtype, D: Device<E>, T: Tape<E, D>> Module<Tensor<S, E, D, T>>
+    for LeakyReLU<E>
+where
+    Tensor<S, E, D, T>: TryPReLU<E>,
+{
+    type Output = Tensor<S, E, D, T>;
+    type Error = <Tensor<S, E, D, T> as HasErr>::Err;
+
+    fn try_forward(&self, input: Tensor<S, E, D, T>) -> Result<Self::Output, Self::Error> {
+        input.try_prelu(self.0)
+    }
+}
+
+pub mod builder {
+    #[derive(Debug, Copy, Clone, Eq, PartialEq)]
+    pub struct PReLU;
+
+    use core::marker::PhantomData;
+
+    use crate::prelude::ConstDim;
+
+    #[derive(Debug, Copy, Clone, Eq, PartialEq)]
+    pub struct PReLU1D<C: ConstDim>(PhantomData<C>);
+}
+
+impl<E: Dtype, D: Device<E>> BuildOnDevice<D, E> for builder::PReLU
+where
+    PReLU<E, D>: BuildModule<D, E>,
+{
+    type Built = PReLU<E, D>;
+    fn try_build_on_device(device: &D) -> Result<Self::Built, <D>::Err> {
+        Self::Built::try_build(device)
+    }
+}
+
+/// Calls [prelu()] with learnable value.
+#[derive(Debug, Clone)]
+pub struct PReLU<E: Dtype, D: Device<E>> {
+    pub a: Tensor<(), E, D>,
+}
+
+impl<E: Dtype, D: Device<E>> NonMutableModule for PReLU<E, D> {}
+
+impl<S: ConstShape, E: Dtype, D: Device<E>, T: Tape<E, D>> Module<Tensor<S, E, D, T>>
+    for PReLU<E, D>
+where
+    Tensor<S, E, D, T>: TryPReLU<Tensor<S, E, D, NoneTape>>,
+{
+    type Output = Tensor<S, E, D, T>;
+    type Error = <Tensor<S, E, D, T> as HasErr>::Err;
+
+    fn try_forward(&self, input: Tensor<S, E, D, T>) -> Result<Self::Output, Self::Error> {
+        input.try_prelu(self.a.retaped().broadcast())
+    }
+}
+
+impl<E: Dtype, D: Device<E>> TensorCollection<E, D> for PReLU<E, D> {
+    type To<E2: Dtype, D2: Device<E2>> = PReLU<E2, D2>;
+
+    fn iter_tensors<V: crate::prelude::ModuleVisitor<Self, E, D>>(
+        visitor: &mut V,
+    ) -> Result<Option<Self::To<V::E2, V::D2>>, V::Err> {
+        visitor.visit_fields(
+            Self::tensor("a", |p| &p.a, |p| &mut p.a, TensorOptions::reset_to_zeros()),
+            |a| PReLU { a },
+        )
+    }
+}
+
+/// Calls [prelu()] with learnable values along second dimension.
+#[derive(Debug, Clone)]
+pub struct PReLU1D<C: ConstDim, E: Dtype, D: Device<E>> {
+    pub a: Tensor<(C,), E, D>,
+}
+
+impl<C: ConstDim, E: Dtype, D: Device<E>> BuildOnDevice<D, E> for builder::PReLU1D<C>
+where
+    PReLU1D<C, E, D>: BuildModule<D, E>,
+{
+    type Built = PReLU1D<C, E, D>;
+    fn try_build_on_device(device: &D) -> Result<Self::Built, <D>::Err> {
+        Self::Built::try_build(device)
+    }
+}
+
+impl<C: ConstDim, E: Dtype, D: Device<E>> NonMutableModule for PReLU1D<C, E, D> {}
+
+impl<C: ConstDim, E: Dtype, D: Device<E>, T: Tape<E, D>> Module<Tensor<(C,), E, D, T>>
+    for PReLU1D<C, E, D>
+where
+    Tensor<(C,), E, D, T>: TryPReLU<Tensor<(C,), E, D, NoneTape>>,
+{
+    type Output = Tensor<(C,), E, D, T>;
+
+    type Error = <Tensor<(C,), E, D, T> as HasErr>::Err;
+
+    fn try_forward(&self, input: Tensor<(C,), E, D, T>) -> Result<Self::Output, Self::Error> {
+        input.try_prelu(self.a.retaped())
+    }
+}
+
+macro_rules! prelu1d {
+    (($($InDims:tt),*), $Axes:ty) => {
+        impl<E: Dtype, D: Device<E>, T: Tape<E, D> + Merge<NoneTape>, $($InDims: ConstDim),*> Module<Tensor<($($InDims),*), E, D, T>> for PReLU1D<C,E, D>
+        where ($($InDims),*): ReduceShapeTo<(C,), $Axes>,
+        Tensor<($($InDims),*), E, D, T>: TryPReLU<Tensor<($($InDims),*), E, D, NoneTape>>,
+        {
+            type Output = Tensor<($($InDims),*), E, D, T>;
+            type Error = <Tensor<($($InDims),*), E, D, T> as HasErr>::Err;
+
+            fn try_forward(&self, input: Tensor<($($InDims),*), E, D, T>) -> Result<Self::Output, Self::Error> {
+                input.try_prelu(self.a.retaped().broadcast())
+            }
+        }
+    };
+}
+
+prelu1d!((B, C), Axis<0>);
+prelu1d!((B, C, M), Axes2<0, 2>);
+prelu1d!((B, C, M, N), Axes3<0, 2, 3>);
+
+impl<C: ConstDim, E: Dtype, D: Device<E>> TensorCollection<E, D> for PReLU1D<C, E, D> {
+    type To<E2: Dtype, D2: Device<E2>> = PReLU1D<C, E2, D2>;
+
+    fn iter_tensors<V: crate::prelude::ModuleVisitor<Self, E, D>>(
+        visitor: &mut V,
+    ) -> Result<Option<Self::To<V::E2, V::D2>>, V::Err> {
+        visitor.visit_fields(
+            Self::tensor("a", |p| &p.a, |p| &mut p.a, TensorOptions::reset_to_zeros()),
+            |a| PReLU1D { a },
+        )
+    }
+}
+
 #[cfg(test)]
 mod tests {
-    use crate::{nn::*, tests::TestDevice};
+    use crate::{
+        nn::*,
+        tests::{assert_close, TestDevice},
+    };
 
     use super::*;
 
@@ -170,4 +325,83 @@ mod tests {
         let r2 = t.softmax::<crate::shapes::Axis<1>>();
         assert_eq!(r1.array(), r2.array());
     }
+
+    #[test]
+    fn test_nn_activations_leaky_relu() {
+        let dev: TestDevice = Default::default();
+
+        let t = dev.tensor([-2.0, -1.0, 0.0, 1.0, 2.0]);
+        let r1 = LeakyReLU(0.05).forward_mut(t.clone());
+        let r2 = t.prelu(dev.tensor([0.05, 0.05, 0.05, 0.05, 0.05]));
+        assert_eq!(r1.array(), r2.array());
+
+        let t = dev.tensor([[-2.0, -1.0, 0.0], [1.0, 2.0, 3.0]]);
+        let r1 = LeakyReLU(0.05).forward_mut(t.clone());
+        let r2 = t.prelu(dev.tensor([[0.05, 0.05, 0.05], [0.05, 0.05, 0.05]]));
+        assert_eq!(r1.array(), r2.array());
+    }
+
+    #[test]
+    fn test_nn_activations_prelu() {
+        let dev: TestDevice = Default::default();
+
+        let t = dev.tensor([-2.0, -1.0, 0.0, 1.0, 2.0]);
+        let r1 = PReLU {
+            a: dev.tensor(0.05),
+        }
+        .forward_mut(t.clone());
+        let r2 = t.prelu(dev.tensor([0.05, 0.05, 0.05, 0.05, 0.05]));
+        assert_eq!(r1.array(), r2.array());
+
+        let t = dev.tensor([[-2.0, -1.0, 0.0], [1.0, 2.0, 3.0]]);
+        let r1 = PReLU {
+            a: dev.tensor(0.05),
+        }
+        .forward_mut(t.clone());
+        let r2 = t.prelu(dev.tensor([[0.05, 0.05, 0.05], [0.05, 0.05, 0.05]]));
+        assert_eq!(r1.array(), r2.array());
+
+        let model = (
+            Tanh,
+            PReLU {
+                a: dev.tensor(0.05),
+            },
+        );
+        let t = dev.tensor([-2.0, -1.0, 0.0, 1.0, 2.0]);
+        let out = model.forward(t);
+        assert_close(
+            &out.array(),
+            &[-0.04820138, -0.03807970, 0.0, 0.76159415, 0.96402758],
+        )
+    }
+
+    #[test]
+    fn test_nn_activations_prelu_1d() {
+        let dev: TestDevice = Default::default();
+
+        let t = dev.tensor([-2.0, -1.0, 0.0, 1.0, 2.0]);
+        let r1 = PReLU1D {
+            a: dev.tensor([0.05, 0.06, 0.07, 0.08, 0.09]),
+        }
+        .forward_mut(t.clone());
+        let r2 = t.prelu(dev.tensor([0.05, 0.06, 0.07, 0.08, 0.09]));
+        assert_eq!(r1.array(), r2.array());
+
+        let t = dev.tensor([[-2.0, -1.0, 0.0], [1.0, 2.0, 3.0]]);
+        let r1 = PReLU1D {
+            a: dev.tensor([0.05, 0.07, 0.09]),
+        }
+        .forward_mut(t.clone());
+        let r2 = t.prelu(dev.tensor([[0.05, 0.07, 0.09], [0.05, 0.07, 0.09]]));
+        assert_eq!(r1.array(), r2.array());
+
+        let mut model = dev.build_module::<(Tanh, builder::PReLU1D<Const<5>>), f32>();
+        let t = dev.tensor([-2.0, -1.0, 0.0, 1.0, 2.0]);
+        model.1.a = dev.tensor([0.05, 0.05, 0.05, 0.05, 0.05]);
+        let out = model.forward(t);
+        assert_close(
+            &out.array(),
+            &[-0.04820138, -0.03807970, 0.0, 0.76159415, 0.96402758],
+        )
+    }
 }
diff --git a/src/nn/mod.rs b/src/nn/mod.rs
index 1b3cd88a8..9f24fa2a0 100644
--- a/src/nn/mod.rs
+++ b/src/nn/mod.rs
@@ -261,12 +261,16 @@ pub mod modules {
     };
     pub use super::unbiased_linear::UnbiasedLinear;
     pub use super::upscale::Upscale2D;
+    #[cfg(feature = "nightly")]
+    pub use super::upscale::Upscale2DBy;
     pub use super::*;
 }
 
 pub mod builders {
     //! Simple specification of network structure, without
     //! worrying about device or dtype.
+    pub use super::activations::builder::PReLU;
+    pub use super::activations::builder::PReLU1D;
     pub use super::activations::*;
     pub use super::add_into::AddInto;
     pub use super::batchnorm1d::builder::BatchNorm1D;
@@ -295,5 +299,7 @@ pub mod builders {
     };
     pub use super::unbiased_linear::builder::UnbiasedLinear;
     pub use super::upscale::Upscale2D;
+    #[cfg(feature = "nightly")]
+    pub use super::upscale::Upscale2DBy;
     pub use super::*;
 }
diff --git a/src/tensor_ops/mod.rs b/src/tensor_ops/mod.rs
index 4c16d56fe..5d1a8bd73 100644
--- a/src/tensor_ops/mod.rs
+++ b/src/tensor_ops/mod.rs
@@ -181,6 +181,7 @@ mod negate;
 mod normalize;
 mod permute_to;
 mod pow;
+mod prelu;
 mod realize_to;
 mod relu;
 mod reshape_to;
@@ -233,6 +234,7 @@ pub use negate::negate;
 pub use normalize::normalize;
 pub use permute_to::PermuteTo;
 pub use pow::{powf, powi};
+pub use prelu::{leakyrelu, prelu, TryPReLU};
 pub use realize_to::RealizeTo;
 pub use relu::relu;
 pub use reshape_to::ReshapeTo;
diff --git a/src/tensor_ops/prelu/mod.rs b/src/tensor_ops/prelu/mod.rs
new file mode 100644
index 000000000..8c10b84a3
--- /dev/null
+++ b/src/tensor_ops/prelu/mod.rs
@@ -0,0 +1,98 @@
+use num_traits::Zero;
+
+use crate::{shapes::*, tensor::*};
+
+use super::{
+    cmp::{LtKernelOp, ScalarCmpKernel},
+    BroadcastTo, ChooseFrom, Device,
+};
+
+/// [Parametric Rectified Linear Unit (PReLU)](https://pytorch.org/docs/stable/generated/torch.nn.PReLU.html). `max(0, t) + a*min(0, t)`
+///
+/// Examples:
+/// ```rust
+/// # use dfdx::prelude::*;
+/// # let dev: Cpu = Default::default();
+/// let t = dev.tensor([-1.0, 0.0, 1.0, 2.0]);
+/// let a = dev.tensor([0.05,0.05,0.05,0.05]);
+/// let r = prelu(t, a);
+/// assert_eq!(r.array(), [-0.05, 0.0, 1.0, 2.0]);
+/// ```
+
+pub fn prelu<S: Shape, E: Dtype, D: Device<E>, T: Tape<E, D> + Merge<R>, R: Default>(
+    lhs: Tensor<S, E, D, T>,
+    rhs: Tensor<S, E, D, R>,
+) -> Tensor<S, E, D, T>
+where
+    Tensor<S, E, D, T>: TryPReLU<Tensor<S, E, D, R>>,
+{
+    lhs.prelu(rhs)
+}
+
+pub fn leakyrelu<S: Shape, E: Dtype, D: Device<E>, T: Tape<E, D>>(
+    lhs: Tensor<S, E, D, T>,
+    rhs: E,
+) -> Tensor<S, E, D, T>
+where
+    Tensor<S, E, D, T>: TryPReLU<E>,
+{
+    lhs.prelu(rhs)
+}
+
+pub trait TryPReLU<T = Self>: HasErr {
+    fn try_prelu(self, rhs: T) -> Result<Self, Self::Err>;
+
+    fn prelu(self, rhs: T) -> Self {
+        self.try_prelu(rhs).unwrap()
+    }
+}
+
+impl<S: Shape, E: Dtype + Zero, D, LhsTape: Tape<E, D>, R> TryPReLU<Tensor<S, E, D, R>>
+    for Tensor<S, E, D, LhsTape>
+where
+    D: Device<E> + ScalarCmpKernel<LtKernelOp, E>,
+    LhsTape: Merge<R>,
+{
+    /// See [prelu]
+    fn try_prelu(self, rhs: Tensor<S, E, D, R>) -> Result<Self, Self::Err> {
+        let scaled = self.with_empty_tape() * rhs;
+        Ok(self.scalar_lt(E::zero()).choose(scaled, self))
+    }
+}
+
+impl<S: Shape, E: Dtype + Zero, D: Device<E> + ScalarCmpKernel<LtKernelOp, E>, T: Tape<E, D>>
+    TryPReLU<E> for Tensor<S, E, D, T>
+{
+    /// See [prelu]
+    fn try_prelu(self, rhs: E) -> Result<Self, Self::Err> {
+        let dev = self.device.clone();
+        let scale = dev.tensor(rhs).retaped::<T>().broadcast_like(self.shape());
+        let scaled = self.with_empty_tape() * scale;
+        Ok(self.scalar_lt(E::zero()).choose(scaled, self))
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use crate::{
+        tensor::*,
+        tensor_ops::{prelu::TryPReLU, *},
+        tests::*,
+    };
+
+    #[test]
+    fn test_prelu() {
+        let dev: TestDevice = Default::default();
+        let x: Tensor<_, TestDtype, _> = dev.tensor([-2.0, -1.0, 0.0, 1.0, 2.0]);
+        let y: Tensor<_, TestDtype, _> = dev.tensor([0.05, 0.05, 0.05, 0.05, 0.05]);
+        let r = x.leaky_trace().prelu(y.clone());
+        assert_eq!(r.array(), [-0.1, -0.05, 0.0, 1.0, 2.0]);
+        // NOTE: call .exp() to make sure we cover cases where .prelu() uses the result's gradient
+        let g = r.exp().mean().backward();
+        assert_close(
+            &g.get(&x).array(),
+            &[0.00904837, 0.00951229, 0.2, 0.54365635, 1.4778112],
+        );
+        assert_close(&g.get(&y).array(), &[-0.3619348, -0.1902458, 0.0, 0.0, 0.0]);
+    }
+}


Test patch: