coreylowman · coreylowman · Apr 5, 2023 · Apr 5, 2023 · Apr 5, 2023 · Apr 5, 2023
diff --git a/src/tensor_ops/mod.rs b/src/tensor_ops/mod.rs
@@ -183,6 +183,7 @@ mod permute_to;
 mod pow;
 mod prelu;
 mod realize_to;
+mod recip;
 mod relu;
 mod reshape_to;
 mod roll;
@@ -236,6 +237,7 @@ pub use permute_to::PermuteTo;
 pub use pow::{powf, powi};
 pub use prelu::{leakyrelu, prelu, TryPReLU};
 pub use realize_to::RealizeTo;
+pub use recip::recip;
 pub use relu::relu;
 pub use reshape_to::ReshapeTo;
 pub use roll::Roll;

diff --git a/src/tensor_ops/recip/cpu_kernel.rs b/src/tensor_ops/recip/cpu_kernel.rs
@@ -0,0 +1,13 @@
+use crate::tensor_ops::cpu_kernels::UnaryDerivative;
+
+impl<F: num_traits::Float> UnaryDerivative<F> for super::RecipKernelOp {
+    const DF_USES_FX: bool = true;
+    #[inline(always)]
+    fn f(&self, x: &F) -> F {
+        x.recip()
+    }
+    #[inline(always)]
+    fn df(&self, fx: &F) -> F {
+        -fx.powi(2)
+    }
+}
diff --git a/src/tensor_ops/recip/cuda_kernel.rs b/src/tensor_ops/recip/cuda_kernel.rs
@@ -0,0 +1,9 @@
+use super::RecipKernelOp;
+use crate::tensor_ops::cuda_kernels::cuda_unary;
+
+unsafe impl cudarc::driver::DeviceRepr for RecipKernelOp {}
+
+const PTX: &str = include_str!(concat!(env!("OUT_DIR"), "/recip.ptx"));
+
+cuda_unary!(RecipKernelOp, f32, PTX, "recip_fwd_f32", "recip_bwd_f32");
+cuda_unary!(RecipKernelOp, f64, PTX, "recip_fwd_f64", "recip_bwd_f64");
diff --git a/src/tensor_ops/recip/mod.rs b/src/tensor_ops/recip/mod.rs
@@ -0,0 +1,55 @@
+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 RecipKernelOp;
+
+/// `1 / x`
+///
+/// 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.recip();
+/// ```
+pub fn recip<S: Shape, E: Dtype, D: UnaryKernel<RecipKernelOp, E>, T: Tape<E, D>>(
+    t: Tensor<S, E, D, T>,
+) -> Tensor<S, E, D, T> {
+    t.recip()
+}
+
+impl<S: Shape, E: Dtype, D: UnaryKernel<RecipKernelOp, E>, T: Tape<E, D>> Tensor<S, E, D, T> {
+    /// See [recip]
+    pub fn recip(self) -> Self {
+        self.try_recip().unwrap()
+    }
+    /// See [recip]
+    pub fn try_recip(self) -> Result<Self, D::Err> {
+        try_unary_op(RecipKernelOp, self)
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use crate::{tensor::*, tensor_ops::*, tests::*};
+
+    #[test]
+    fn test_recip() {
+        let dev: TestDevice = Default::default();
+        let x: Tensor<_, TestDtype, _> = dev.tensor([-2.0, -1.0, 0.0, 1.0, 2.0]);
+        let r = x.leaky_trace().recip();
+        assert_close(&r.array(), &[-0.5, -1.0, TestDtype::INFINITY, 1.0, 0.5]);
+        let g = r.mean().backward();
+        assert_close(
+            &g.get(&x).array(),
+            &[-0.05, -0.2, TestDtype::NEG_INFINITY, -0.2, -0.05],
+        );
+    }
+}
diff --git a/src/tensor_ops/recip/recip.cu b/src/tensor_ops/recip/recip.cu
@@ -0,0 +1,15 @@
+#include "unary_op_macros.cuh"
+
+struct RecipKernelOp {};
+
+UNARY_OP(
+    float, recip_fwd_f32, recip_bwd_f32, RecipKernelOp,
+    1 / x,
+    -y * y
+)
+
+UNARY_OP(
+    double, recip_fwd_f64, recip_bwd_f64, RecipKernelOp,
+    1 / x,
+    -y * y
+)
diff --git a/src/tensor_ops/utilities/device.rs b/src/tensor_ops/utilities/device.rs
@@ -87,6 +87,7 @@ pub trait Device<E: Dtype>:
     + UnaryKernel<super::super::tanh::TanhKernelOp, E>
     + UnaryKernel<super::super::pow::PowfKernelOp<E>, E>
     + UnaryKernel<super::super::pow::PowiKernelOp, E>
+    + UnaryKernel<super::super::recip::RecipKernelOp, E>
 
     // to_dtype
     + super::super::to_dtype::ToDtypeKernel<f32, E>