coreylowman · coreylowman · Apr 4, 2023 · Apr 4, 2023 · Apr 4, 2023 · Apr 4, 2023
diff --git a/src/tensor_ops/abs/cpu_kernel.rs b/src/tensor_ops/abs/cpu_kernel.rs
@@ -2,6 +2,7 @@ use crate::tensor_ops::cpu_kernels::UnaryDerivative;
 use num_traits::Float;
 
 impl<F: Float> UnaryDerivative<F> for super::AbsKernelOp {
+    const DF_USES_FX: bool = false;
     #[inline(always)]
     fn f(&self, x: &F) -> F {
         x.abs()

diff --git a/src/tensor_ops/add/cpu_kernel.rs b/src/tensor_ops/add/cpu_kernel.rs
@@ -17,9 +17,12 @@ impl<F: Float> BinaryDerivative<F> for super::BinaryAddKernelOp {
 }
 
 impl<F: Float> UnaryDerivative<F> for super::ScalarAddKernelOp<F> {
+    const DF_USES_FX: bool = false;
+    #[inline(always)]
     fn f(&self, &x: &F) -> F {
         x + self.scalar
     }
+    #[inline(always)]
     fn df(&self, _: &F) -> F {
         F::one()
     }

diff --git a/src/tensor_ops/clamp/cpu_kernel.rs b/src/tensor_ops/clamp/cpu_kernel.rs
@@ -2,6 +2,7 @@ use crate::tensor_ops::cpu_kernels::UnaryDerivative;
 use num_traits::{clamp, Float};
 
 impl<F: Float + PartialOrd> UnaryDerivative<F> for super::ClampKernelOp<F> {
+    const DF_USES_FX: bool = false;
     #[inline(always)]
     fn f(&self, &x: &F) -> F {
         clamp(x, self.min, self.max)

diff --git a/src/tensor_ops/cos/cpu_kernel.rs b/src/tensor_ops/cos/cpu_kernel.rs
@@ -2,6 +2,7 @@ use crate::tensor_ops::cpu_kernels::UnaryDerivative;
 use num_traits::Float;
 
 impl<F: Float> UnaryDerivative<F> for super::CosKernelOp {
+    const DF_USES_FX: bool = false;
     #[inline(always)]
     fn f(&self, x: &F) -> F {
         x.cos()

diff --git a/src/tensor_ops/div/cpu_kernel.rs b/src/tensor_ops/div/cpu_kernel.rs
@@ -2,9 +2,12 @@ use crate::tensor_ops::cpu_kernels::{BinaryDerivative, UnaryDerivative};
 use num_traits::Float;
 
 impl<F: Float> UnaryDerivative<F> for super::ScalarDivKernelOp<F> {
+    const DF_USES_FX: bool = false;
+    #[inline(always)]
     fn f(&self, &x: &F) -> F {
         x / self.scalar
     }
+    #[inline(always)]
     fn df(&self, _: &F) -> F {
         F::one() / self.scalar
     }

diff --git a/src/tensor_ops/exp/cpu_kernel.rs b/src/tensor_ops/exp/cpu_kernel.rs
@@ -2,12 +2,13 @@ use crate::tensor_ops::cpu_kernels::UnaryDerivative;
 use num_traits::Float;
 
 impl<F: Float> UnaryDerivative<F> for super::ExpKernelOp {
+    const DF_USES_FX: bool = true;
     #[inline(always)]
     fn f(&self, x: &F) -> F {
         x.exp()
     }
     #[inline(always)]
-    fn df(&self, x: &F) -> F {
-        x.exp()
+    fn df(&self, &fx: &F) -> F {
+        fx
     }
 }
diff --git a/src/tensor_ops/exp/exp.cu b/src/tensor_ops/exp/exp.cu
@@ -4,9 +4,9 @@ struct ExpKernelOp {};
 
 UNARY_OP(float, exp_fwd_f32, exp_bwd_f32, ExpKernelOp,
         expf(x),
-        expf(x))
+        y)
 
 UNARY_OP(double, exp_fwd_f64, exp_bwd_f64, ExpKernelOp,
         exp(x),
-        exp(x))
+        y)
 
diff --git a/src/tensor_ops/gelu/cpu_kernel.rs b/src/tensor_ops/gelu/cpu_kernel.rs
@@ -2,6 +2,7 @@ use crate::tensor_ops::cpu_kernels::UnaryDerivative;
 use num_traits::{Float, FloatConst};
 
 impl<F: Float + FloatConst> UnaryDerivative<F> for super::GeLUKernelOp {
+    const DF_USES_FX: bool = false;
     #[inline(always)]
     fn f(&self, &x: &F) -> F {
         let alpha = x + F::from(0.044715).unwrap() * x.powi(3);

diff --git a/src/tensor_ops/ln/cpu_kernel.rs b/src/tensor_ops/ln/cpu_kernel.rs
@@ -2,6 +2,7 @@ use crate::tensor_ops::cpu_kernels::UnaryDerivative;
 use num_traits::Float;
 
 impl<F: Float> UnaryDerivative<F> for super::LnKernelOp {
+    const DF_USES_FX: bool = false;
     #[inline(always)]
     fn f(&self, x: &F) -> F {
         x.ln()

diff --git a/src/tensor_ops/mul/cpu_kernel.rs b/src/tensor_ops/mul/cpu_kernel.rs
@@ -3,9 +3,12 @@ use crate::tensor_ops::cpu_kernels::{BinaryDerivative, UnaryDerivative};
 use num_traits::Float;
 
 impl<F: Float> UnaryDerivative<F> for super::ScalarMulKernelOp<F> {
+    const DF_USES_FX: bool = false;
+    #[inline(always)]
     fn f(&self, &x: &F) -> F {
         x * self.scalar
     }
+    #[inline(always)]
     fn df(&self, _: &F) -> F {
         self.scalar
     }

diff --git a/src/tensor_ops/nans_to/cpu_kernel.rs b/src/tensor_ops/nans_to/cpu_kernel.rs
@@ -1,6 +1,7 @@
 use crate::tensor_ops::cpu_kernels::UnaryDerivative;
 
 impl<F: num_traits::Float> UnaryDerivative<F> for super::NansToKernelOp<F> {
+    const DF_USES_FX: bool = false;
     #[inline(always)]
     fn f(&self, x: &F) -> F {
         if x.is_nan() {

diff --git a/src/tensor_ops/negate/cpu_kernel.rs b/src/tensor_ops/negate/cpu_kernel.rs
@@ -1,6 +1,7 @@
 use crate::tensor_ops::cpu_kernels::UnaryDerivative;
 
 impl<F: num_traits::Float> UnaryDerivative<F> for super::NegateKernelOp {
+    const DF_USES_FX: bool = false;
     #[inline(always)]
     fn f(&self, x: &F) -> F {
         x.neg()

diff --git a/src/tensor_ops/pow/cpu_kernel.rs b/src/tensor_ops/pow/cpu_kernel.rs
@@ -1,6 +1,7 @@
 use crate::tensor_ops::cpu_kernels::UnaryDerivative;
 
 impl<F: num_traits::Float> UnaryDerivative<F> for super::PowiKernelOp {
+    const DF_USES_FX: bool = false;
     #[inline(always)]
     fn f(&self, x: &F) -> F {
         x.powi(self.0)
@@ -12,6 +13,7 @@ impl<F: num_traits::Float> UnaryDerivative<F> for super::PowiKernelOp {
 }
 
 impl<F: num_traits::Float> UnaryDerivative<F> for super::PowfKernelOp<F> {
+    const DF_USES_FX: bool = false;
     #[inline(always)]
     fn f(&self, x: &F) -> F {
         x.powf(self.0)

diff --git a/src/tensor_ops/pow/cuda_kernel.rs b/src/tensor_ops/pow/cuda_kernel.rs
@@ -30,12 +30,14 @@ where
         op: super::PowiKernelOp,
         inp: &Tensor<S, E, Self>,
         grad_inp: &mut Self::Vec<E>,
+        out: &Tensor<S, E, Self>,
         grad_out: &Self::Vec<E>,
     ) -> Result<(), Self::Err> {
         self.backward(
             super::PowfKernelOp(E::from_i32(op.0).unwrap()),
             inp,
             grad_inp,
+            out,
             grad_out,
         )
     }

diff --git a/src/tensor_ops/relu/cpu_kernel.rs b/src/tensor_ops/relu/cpu_kernel.rs
@@ -1,6 +1,7 @@
 use crate::tensor_ops::cpu_kernels::UnaryDerivative;
 
 impl<F: num_traits::Float> UnaryDerivative<F> for super::ReLUKernelOp {
+    const DF_USES_FX: bool = false;
     #[inline(always)]
     fn f(&self, x: &F) -> F {
         x.max(F::zero())

diff --git a/src/tensor_ops/sigmoid/cpu_kernel.rs b/src/tensor_ops/sigmoid/cpu_kernel.rs
@@ -1,13 +1,13 @@
 use crate::tensor_ops::cpu_kernels::UnaryDerivative;
 
 impl<F: num_traits::Float> UnaryDerivative<F> for super::SigmoidKernelOp {
+    const DF_USES_FX: bool = true;
     #[inline(always)]
     fn f(&self, x: &F) -> F {
         F::one() / (F::one() + x.neg().exp())
     }
     #[inline(always)]
-    fn df(&self, x: &F) -> F {
-        let fx = self.f(x);
+    fn df(&self, &fx: &F) -> F {
         fx * (F::one() - fx)
     }
 }
diff --git a/src/tensor_ops/sigmoid/sigmoid.cu b/src/tensor_ops/sigmoid/sigmoid.cu
@@ -7,9 +7,9 @@ struct SigmoidKernelOp {};
 
 UNARY_OP(float, sigmoid_fwd_f32, sigmoid_bwd_f32, SigmoidKernelOp,
         SIGMOID_f32(x),
-        SIGMOID_f32(x) * (1.0 - SIGMOID_f32(x)))
+        y * (1.0 - y))
 
 UNARY_OP(double, sigmoid_fwd_f64, sigmoid_bwd_f64, SigmoidKernelOp,
         SIGMOID_f64(x),
-        SIGMOID_f64(x) * (1.0 - SIGMOID_f64(x)))
+        y * (1.0 - y))
 
diff --git a/src/tensor_ops/sin/cpu_kernel.rs b/src/tensor_ops/sin/cpu_kernel.rs
@@ -1,6 +1,7 @@
 use crate::tensor_ops::cpu_kernels::UnaryDerivative;
 
 impl<F: num_traits::Float> UnaryDerivative<F> for super::SinKernelOp {
+    const DF_USES_FX: bool = false;
     #[inline(always)]
     fn f(&self, x: &F) -> F {
         x.sin()

diff --git a/src/tensor_ops/sqrt/cpu_kernel.rs b/src/tensor_ops/sqrt/cpu_kernel.rs
@@ -1,12 +1,13 @@
 use crate::tensor_ops::cpu_kernels::UnaryDerivative;
 
 impl<F: num_traits::Float> UnaryDerivative<F> for super::SqrtKernelOp {
+    const DF_USES_FX: bool = true;
     #[inline(always)]
     fn f(&self, x: &F) -> F {
         x.sqrt()
     }
     #[inline(always)]
-    fn df(&self, x: &F) -> F {
-        F::from(0.5).unwrap() / x.sqrt()
+    fn df(&self, &fx: &F) -> F {
+        (fx + fx).recip()
     }
 }
diff --git a/src/tensor_ops/sqrt/sqrt.cu b/src/tensor_ops/sqrt/sqrt.cu
@@ -4,9 +4,9 @@ struct SqrtKernelOp {};
 
 UNARY_OP(float, sqrt_fwd_f32, sqrt_bwd_f32, SqrtKernelOp,
         sqrtf(x),
-        0.5 / sqrtf(x))
+        1 / (y + y))
 
 UNARY_OP(double, sqrt_fwd_f64, sqrt_bwd_f64, SqrtKernelOp,
         sqrt(x),
-        0.5 / sqrt(x))
+        1 / (y + y))
 
diff --git a/src/tensor_ops/square/cpu_kernel.rs b/src/tensor_ops/square/cpu_kernel.rs
@@ -1,6 +1,7 @@
 use crate::tensor_ops::cpu_kernels::UnaryDerivative;
 
 impl<F: num_traits::Float> UnaryDerivative<F> for super::SquareKernelOp {
+    const DF_USES_FX: bool = false;
     #[inline(always)]
     fn f(&self, x: &F) -> F {
         x.powi(2)

diff --git a/src/tensor_ops/sub/cpu_kernel.rs b/src/tensor_ops/sub/cpu_kernel.rs
@@ -1,9 +1,12 @@
 use crate::tensor_ops::cpu_kernels::{BinaryDerivative, UnaryDerivative};
 
 impl<F: num_traits::Float> UnaryDerivative<F> for super::ScalarSubKernelOp<F> {
+    const DF_USES_FX: bool = false;
+    #[inline(always)]
     fn f(&self, &x: &F) -> F {
         x - self.scalar
     }
+    #[inline(always)]
     fn df(&self, _: &F) -> F {
         F::one()
     }

diff --git a/src/tensor_ops/tanh/cpu_kernel.rs b/src/tensor_ops/tanh/cpu_kernel.rs
@@ -1,12 +1,13 @@
 use crate::tensor_ops::cpu_kernels::UnaryDerivative;
 
 impl<F: num_traits::Float> UnaryDerivative<F> for super::TanhKernelOp {
+    const DF_USES_FX: bool = true;
     #[inline(always)]
     fn f(&self, x: &F) -> F {
         x.tanh()
     }
     #[inline(always)]
-    fn df(&self, x: &F) -> F {
-        F::one() - x.tanh().powi(2)
+    fn df(&self, fx: &F) -> F {
+        F::one() - fx.powi(2)
     }
 }
diff --git a/src/tensor_ops/tanh/tanh.cu b/src/tensor_ops/tanh/tanh.cu
@@ -4,9 +4,9 @@ struct TanhKernelOp {};
 
 UNARY_OP(float, tanh_fwd_f32, tanh_bwd_f32, TanhKernelOp,
         tanhf(x),
-        1 - tanhf(x) * tanhf(x))
+        1 - y * y)
 
 UNARY_OP(double, tanh_fwd_f64, tanh_bwd_f64, TanhKernelOp,
         tanh(x),
-        1 - tanh(x) * tanh(x))
+        1 - y * y)
 
diff --git a/src/tensor_ops/utilities/cpu_kernels.rs b/src/tensor_ops/utilities/cpu_kernels.rs
@@ -8,7 +8,13 @@ use crate::{
 };
 
 pub trait UnaryDerivative<E> {
+    /// Whether the [UnaryDerivative::df] function can re-use the output
+    /// from [UnaryDerivative::f].
+    const DF_USES_FX: bool;
     fn f(&self, x: &E) -> E;
+
+    /// Receives `f(x)` if [UnaryDerivative::DF_USES_FX] is true,
+    /// otherwise `x`.
     fn df(&self, x: &E) -> E;
 }
 
@@ -45,12 +51,19 @@ impl<E: Dtype, Op: UnaryDerivative<E>> UnaryKernel<Op, E> for Cpu {
         op: Op,
         inp: &Tensor<S, E, Self>,
         grad_inp: &mut Self::Vec<E>,
+        out: &Tensor<S, E, Self>,
         grad_out: &Self::Vec<E>,
     ) -> Result<(), Self::Err> {
         debug_assert_eq!(grad_inp.len(), grad_out.len());
         debug_assert_eq!(inp.data.len(), grad_out.len());
-        for (i, x) in grad_inp.iter_mut().enumerate() {
-            *x += op.df(&inp.data[i]) * grad_out[i];
+        if Op::DF_USES_FX {
+            for (i, x) in grad_inp.iter_mut().enumerate() {
+                *x += op.df(&out.data[i]) * grad_out[i];
+            }
+        } else {
+            for (i, x) in grad_inp.iter_mut().enumerate() {
+                *x += op.df(&inp.data[i]) * grad_out[i];
+            }
         }
         Ok(())
     }

diff --git a/src/tensor_ops/utilities/cuda_kernels.rs b/src/tensor_ops/utilities/cuda_kernels.rs
@@ -69,12 +69,20 @@ impl<E: Dtype, K: UnaryOpCudaKernel<E> + DeviceRepr> UnaryKernel<K, E> for Cuda
         op: K,
         inp: &Tensor<S, E, Self>,
         grad_inp: &mut Self::Vec<E>,
+        out: &Tensor<S, E, Self>,
         grad_out: &Self::Vec<E>,
     ) -> Result<(), Self::Err> {
         let bwd_fn = self.dev.get_func(K::MODULE_NAME, K::BWD_FN_NAME).unwrap();
         let numel = inp.data.len();
         let cfg = launch_cfg(numel as u32);
-        let params = (op, numel, inp.data.as_ref(), grad_inp, grad_out);
+        let params = (
+            op,
+            numel,
+            inp.data.as_ref(),
+            grad_inp,
+            out.data.as_ref(),
+            grad_out,
+        );
         unsafe { bwd_fn.launch(cfg, params) }?;
         Ok(())
     }

diff --git a/src/tensor_ops/utilities/ops.rs b/src/tensor_ops/utilities/ops.rs
@@ -14,6 +14,7 @@ pub trait UnaryKernel<Op, E: Dtype>: DeviceStorage {
         op: Op,
         inp: &Tensor<S, E, Self>,
         grad_inp: &mut Self::Vec<E>,
+        out: &Tensor<S, E, Self>,
         grad_out: &Self::Vec<E>,
     ) -> Result<(), Self::Err>;
 }
@@ -54,8 +55,8 @@ pub(crate) fn try_unary_op<
     tape.try_alloc_grad(&out)?;
     tape.add_backward_op(move |grads| {
         let (grad_inp, grad_out) = grads.mut_and_ref(&inp, &phantom_out);
-        inp.device.backward(op, &inp, grad_inp, grad_out)?;
-        Ok(())
+        inp.device
+            .backward(op, &inp, grad_inp, &phantom_out, grad_out)
     });
     Ok(out.put_tape(tape))
 }
@@ -84,8 +85,7 @@ pub(crate) fn try_binary_op<
     tape.add_backward_op(move |grads| {
         let (grad_lhs, grad_rhs, grad_out) = grads.muts_and_ref(&lhs, &rhs, &phantom_out);
         lhs.device
-            .backward(op, &lhs, grad_lhs, &rhs, grad_rhs, grad_out)?;
-        Ok(())
+            .backward(op, &lhs, grad_lhs, &rhs, grad_rhs, grad_out)
     });
     Ok(out.put_tape(tape))
 }
diff --git a/src/tensor_ops/utilities/unary_op_macros.cuh b/src/tensor_ops/utilities/unary_op_macros.cuh
@@ -18,6 +18,7 @@ extern "C" __global__ void BACKWARD( \
     const size_t numel, \
     const TYPENAME *inp, \
     TYPENAME *grad_inp, \
+    const TYPENAME *out, \
     const TYPENAME *grad_out \
 ) { \
     unsigned int i = blockIdx.x * blockDim.x + threadIdx.x; \
@@ -26,6 +27,7 @@ extern "C" __global__ void BACKWARD( \
     } \
     \
     TYPENAME x = inp[i]; \
+    TYPENAME y = out[i]; \
     TYPENAME dx; \
     DERIVATIVE \
     grad_inp[i] += dx * grad_out[i]; \