diff --git a/include/aidge/backend/cpu/operator/HeavisideImpl_kernels.hpp b/include/aidge/backend/cpu/operator/HeavisideImpl_kernels.hpp index 06d7fff8776d342da0467ad7f9d6759f45202151..0bbbddee1040224fd9d0c2658f97c57e3956ca48 100644 --- a/include/aidge/backend/cpu/operator/HeavisideImpl_kernels.hpp +++ b/include/aidge/backend/cpu/operator/HeavisideImpl_kernels.hpp @@ -19,7 +19,6 @@ #include "aidge/backend/cpu/operator/HeavisideImpl.hpp" #include "aidge/utils/ErrorHandling.hpp" - namespace Aidge { template <class I, class O> @@ -35,6 +34,24 @@ void HeavisideImplCpuForwardKernel(std::size_t inputLength, } } + +// Surrogate Gradient +template <class O, class GO, class GI> +void HeavisideImplCpuBackwardKernel(std::size_t inputLength, + const void* output_, + const void* grad_output_, + void* grad_input_) { + + const O* output = static_cast<const O*>(output_); + const GO* grad_output = static_cast<const GO*>(grad_output_); + GI* grad_input = static_cast<GI*>(grad_input_); + + for (size_t i = 0; i < inputLength; ++i) { + // dx = dy * (1/PI) * (1 / (1 + (PI * x)^2)) + grad_input[i] = (1 / M_PI) * grad_output[i] * static_cast<O>(1.0 / (1.0 + output[i] * output[i])); + } +} + // Kernels registration to implementation entry point REGISTRAR(HeavisideImplCpu, {DataType::Float32}, diff --git a/src/operator/HeavisideImpl.cpp b/src/operator/HeavisideImpl.cpp index 56ceb9b0b474d416f25d77b533373d4b193532b8..2ead2978ef1c17ec38e2d12df6b1c78f6c894964 100644 --- a/src/operator/HeavisideImpl.cpp +++ b/src/operator/HeavisideImpl.cpp @@ -32,6 +32,23 @@ template <> void Aidge::HeavisideImplCpu::forward() { op_.value()); } -template <> void Aidge::HeavisideImplCpu::backward() { +template <> +void Aidge::HeavisideImplCpu::backward() { AIDGE_THROW_OR_ABORT(std::runtime_error, "Heaviside backward not implemented yet"); + + // TODO: The following lines are assuming that the surrogate gradient is Atan + // remove that assumption by providing an attribute to Heaviside, + // allowing to choose between different surrogate gradients. + + // const Heavisde_Op& op_ = dynamic_cast<const Heavisie_Op &>(mOp); + + + + // ! backward of hs = forward of atan + //const auto impl = Registrar<HeavisideImplCpu>::create(getBestMatch(getRequiredSpec())); + // std::shared_ptr<Tensor> in0 = op_.getInput(0); + // std::shared_ptr<Tensor> out0 = op_.getOutput(0); + + //impl.forward() } + diff --git a/unit_tests/operator/Test_HeavisideImpl.cpp b/unit_tests/operator/Test_HeavisideImpl.cpp index 4cbdf1a0e29f8670e45897236374726dac62bb43..a0142513d20d4df474d69582e89c8648a06fa340 100644 --- a/unit_tests/operator/Test_HeavisideImpl.cpp +++ b/unit_tests/operator/Test_HeavisideImpl.cpp @@ -11,6 +11,7 @@ #include "aidge/backend/cpu/operator/HeavisideImpl_kernels.hpp" +#include <aidge/operator/Memorize.hpp> #include <memory> #include <cstdlib> #include <random> @@ -22,6 +23,8 @@ #include "aidge/graph/Node.hpp" #include "aidge/utils/TensorUtils.hpp" +#include "aidge/operator/Add.hpp" + namespace Aidge { @@ -95,4 +98,12 @@ TEST_CASE("[cpu/operator] Heaviside(forward)", "[Heaviside][CPU]") { REQUIRE(approxEq<float>(*(op->getOutput(0)), *T1)); } } + +TEST_CASE("[cpu/operator] Heaviside(backward)", "[Heaviside][CPU]") { + + auto add = Add(); + auto mem = Memorize(2); + auto hs = Heaviside(1); +} + }