Sigmoid with residue

src/nn/activations.rs

@@ -36,6 +36,7 @@ activation_impls!(Cos, try_cos, #[doc="Calls [cos()]."]);
 activation_impls!(Ln, try_ln, #[doc="Calls [ln()]."]);
 activation_impls!(Exp, try_exp, #[doc="Calls [exp()]."]);
 activation_impls!(Sigmoid, try_sigmoid, #[doc="Calls [sigmoid()]."]);
+activation_impls!(Sigmoidr, try_sigmoidr, #[doc="Calls [sigmoidr()]."]);
 activation_impls!(Tanh, try_tanh, #[doc="Calls [tanh()]."]);
 activation_impls!(Square, try_square, #[doc="Calls [square()]."]);
 activation_impls!(Sqrt, try_sqrt, #[doc="Calls [sqrt()]."]);

src/tensor_ops/mod.rs

@@ -195,6 +195,7 @@ mod roll;
 mod select_and_gather;
 mod sgd;
 mod sigmoid;
+mod sigmoidr;
 mod sin;
 mod slice;
 mod softmax;
@@ -259,6 +260,7 @@ pub use roll::Roll;
 pub use select_and_gather::{GatherTo, SelectTo};
 pub use sgd::SgdConfig;
 pub use sigmoid::sigmoid;
+pub use sigmoidr::sigmoidr;
 pub use sin::sin;
 pub use slice::slice;
 pub use softmax::softmax;

src/tensor_ops/sigmoidr/cpu_kernel.rs

@@ -0,0 +1,15 @@
+use crate::tensor_ops::cpu_kernels::UnaryDerivative;
+
+impl<F: num_traits::Float> UnaryDerivative<F> for super::SigmoidrKernelOp {
+    const DF_USES_FX: bool = true;
+    const HAS_CONST_DF: bool = false;
+    #[inline(always)]
+    fn f(&self, x: &F) -> F {
+        F::one() / (F::one() + x.neg().exp())
+    }
+    #[inline(always)]
+    fn df(&self, &fx: &F) -> F {
+        let d = fx * (F::one() - fx);
+        F::max(d, F::from(0.0000001).unwrap())
+    }
+}

src/tensor_ops/sigmoidr/cuda_kernel.rs

@@ -0,0 +1,15 @@
+use super::SigmoidrKernelOp as Sigmoidr;
+#[allow(unused_imports)]
+use crate::dtypes::*;
+use crate::tensor_ops::cuda_kernels::cuda_unary;
+
+unsafe impl cudarc::driver::DeviceRepr for Sigmoidr {}
+
+const PTX: &str = include_str!(concat!(env!("OUT_DIR"), "/sigmoidr.ptx"));
+
+#[cfg(feature = "f16")]
+cuda_unary!(df(f(x)) Sigmoidr, f16, PTX, "sigmoidr_fwd_f16", "sigmoidr_bwd_f16");
+#[cfg(feature = "f16")]
+cuda_unary!(df(f(x)) Sigmoidr, AMP<f16>, PTX, "sigmoidr_fwd_f16", "sigmoidr_bwd_f16");
+cuda_unary!(df(f(x)) Sigmoidr, f32, PTX, "sigmoidr_fwd_f32", "sigmoidr_bwd_f32");
+cuda_unary!(df(f(x)) Sigmoidr, f64, PTX, "sigmoidr_fwd_f64", "sigmoidr_bwd_f64");

src/tensor_ops/sigmoidr/mod.rs

@@ -0,0 +1,59 @@
+mod cpu_kernel;
+
+#[cfg(feature = "cuda")]
+mod cuda_kernel;
+
+use super::ops::{try_unary_op, UnaryKernel};
+use crate::{shapes::*, tensor::*};
+
+#[repr(C)]
+#[derive(Debug, Default, Copy, Clone)]
+pub struct SigmoidrKernelOp;
+
+/// [Sigmoid](https://en.wikipedia.org/wiki/Sigmoid_function). `1 / (1 + exp(-t))`.
+/// Basically the same as sigmoid but will always return non-zero gradients.
+/// The derivative is `sigmoid(t) * (1.0 - sigmoid(t))`.
+///
+/// Examples:
+/// ```rust
+/// # use dfdx::prelude::*;
+/// # let dev: Cpu = Default::default();
+/// let t = dev.tensor([-1.0, 0.0, 1.0, 2.0]);
+/// let r = t.sigmoid();
+/// ```
+pub fn sigmoidr<S: Shape, E: Dtype, D: UnaryKernel<SigmoidrKernelOp, E>, T: Tape<E, D>>(
+    t: Tensor<S, E, D, T>,
+) -> Tensor<S, E, D, T> {
+    t.sigmoidr()
+}
+
+impl<S: Shape, E: Dtype, D: UnaryKernel<SigmoidrKernelOp, E>, T: Tape<E, D>> Tensor<S, E, D, T> {
+    /// See [sigmoidr]
+    pub fn sigmoidr(self) -> Self {
+        self.try_sigmoidr().unwrap()
+    }
+    /// See [sigmoidr]
+    pub fn try_sigmoidr(self) -> Result<Self, D::Err> {
+        try_unary_op(SigmoidrKernelOp, self)
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use crate::{tensor::*, tensor_ops::*, tests::*};
+
+    #[test]
+    fn test_sigmoidr() {
+        let dev: TestDevice = Default::default();
+        let x = dev
+            .tensor([-2.0, -1.0, 0.0, 1.0, -TestDtype::INFINITY])
+            .to_dtype::<TestDtype>();
+        let r = x.leaky_trace().sigmoidr();
+        assert_close_to_literal!(r, [0.11920292, 0.26894143, 0.5, 0.7310586, 0.0]);
+        let g = r.mean().backward();
+        assert_close_to_literal!(
+            g.get(&x),
+            [0.020998716, 0.039322387, 0.05, 0.039322387, 0.00000002]
+        );
+    }
+}

src/tensor_ops/sigmoidr/sigmoidr.cu

@@ -0,0 +1,29 @@
+#include "unary_op_macros.cuh"
+
+struct SigmoidrKernelOp {};
+
+template<typename T>
+__device__ __forceinline__ T sigmoidr_fwd(T x) {
+    T one = 1.0;
+    return one / (one + expg(-x));
+}
+
+template<typename T>
+__device__ __forceinline__ T sigmoidr_bwd(T y) {
+    T one = 1.0;
+  T d = y * (one - y);
+    return max(d, 0.0000001);
+}
+
+UNARY_OP(__half, sigmoidr_fwd_f16, sigmoidr_bwd_f16, SigmoidrKernelOp,
+        sigmoidr_fwd(x),
+        sigmoidr_bwd(y))
+
+UNARY_OP(float, sigmoidr_fwd_f32, sigmoidr_bwd_f32, SigmoidrKernelOp,
+        sigmoidr_fwd(x),
+        sigmoidr_bwd(y))
+
+UNARY_OP(double, sigmoidr_fwd_f64, sigmoidr_bwd_f64, SigmoidrKernelOp,
+        sigmoidr_fwd(x),
+        sigmoidr_bwd(y))
+

src/tensor_ops/utilities/device.rs

@@ -89,6 +89,7 @@ pub trait Device<E: Dtype>:
     + UnaryKernel<super::super::fast_gelu::FastGeLUKernelOp, E>
     + UnaryKernel<super::super::accurate_gelu::AccurateGeLUKernelOp, E>
     + UnaryKernel<super::super::sigmoid::SigmoidKernelOp, E>
+    + UnaryKernel<super::super::sigmoidr::SigmoidrKernelOp, E>
     + UnaryKernel<super::super::sin::SinKernelOp, E>
     + UnaryKernel<super::super::sqrt::SqrtKernelOp, E>
     + UnaryKernel<super::super::square::SquareKernelOp, E>