Add pushforwards for the Kokkos:: math functions.

Kokkos ships its own math functions (Kokkos::cos, Kokkos::sqrt, ...) so
a kernel differentiates identically on host and device. clad has no
registered derivatives for them, so it descends into the wrapper's body
and trips an internal assertion instead of differentiating the call.

Register pushforwards for cos, sin, sqrt, and fabs in KokkosBuiltins.h,
mirroring how BuiltinDerivatives.h handles the std:: math functions:
clad now emits a call to the known derivative rather than cloning the
wrapper. Reverse mode is synthesized from the same pushforward, so no
pullback is added. The change is header-only: dispatch to
clad::custom_derivatives::Kokkos::* is the path the View pushforwards in
this header already use.

Author: vgvassilev

include/clad/Differentiator/KokkosBuiltins.h

@@ -322,6 +322,31 @@ void operator_call_pullback(const ::Kokkos::View<DataType, ViewParams...>* v,
 
 /// Kokkos functions (view utils)
 namespace Kokkos {
+/// Pushforwards for the Kokkos math functions -- without a registered
+/// derivative clad differentiates the wrapper body itself and asserts.
+/// Pushforward only, as std:: math is in BuiltinDerivatives.h: reverse mode is
+/// synthesized from it, so a hand-written _pullback would be rejected.
+template <typename T, typename dT>
+KOKKOS_INLINE_FUNCTION clad::ValueAndPushforward<T, dT>
+cos_pushforward(T x, dT d_x) {
+  return {::Kokkos::cos(x), (-1) * ::Kokkos::sin(x) * d_x};
+}
+template <typename T, typename dT>
+KOKKOS_INLINE_FUNCTION clad::ValueAndPushforward<T, dT>
+sin_pushforward(T x, dT d_x) {
+  return {::Kokkos::sin(x), ::Kokkos::cos(x) * d_x};
+}
+template <typename T, typename dT>
+KOKKOS_INLINE_FUNCTION clad::ValueAndPushforward<T, dT>
+sqrt_pushforward(T x, dT d_x) {
+  return {::Kokkos::sqrt(x), d_x / (((T)2) * ::Kokkos::sqrt(x))};
+}
+template <typename T, typename dT>
+KOKKOS_INLINE_FUNCTION clad::ValueAndPushforward<T, dT>
+fabs_pushforward(T x, dT d_x) {
+  return {::Kokkos::fabs(x), (x < 0) ? ((-1) * d_x) : d_x};
+}
+
 template <typename View1, typename View2, typename T>
 inline void deep_copy_pushforward(const View1& dst, const View2& src, T param,
                                   const View1& d_dst, const View2& d_src,

unittests/Kokkos/CMakeLists.txt

@@ -4,6 +4,7 @@ add_clad_unittest(KokkosTests
   ViewBasics.cpp
   ParallelReduce.cpp
   ParallelFor.cpp
+  MathFunctions.cpp
   )
 
 # If llvm does not require rtti, kokkos does.

unittests/Kokkos/MathFunctions.cpp

@@ -0,0 +1,45 @@
+// Pins the Kokkos math pushforwards (cos/sin/sqrt/fabs, KokkosBuiltins.h)
+// against a finite-difference tangent, forward and reverse.
+
+#include "TestUtils.h"
+#include "clad/Differentiator/Differentiator.h"
+#include "clad/Differentiator/KokkosBuiltins.h"
+#include <Kokkos_Core.hpp>
+#include "gtest/gtest.h"
+
+// Exercises all four functions in one expression. sqrt(fabs(x)) keeps the test
+// valid for both signs of x (the fabs kink is only at 0, which we avoid).
+double math_fn(double x, double y) {
+  return Kokkos::sqrt(Kokkos::fabs(x)) * Kokkos::cos(y) + Kokkos::sin(x * y);
+}
+
+TEST(KokkosMath, Forward) {
+  const double eps = 1e-6, tol = 1e-5;
+  auto fn_x = clad::differentiate(math_fn, "x");
+  auto fn_y = clad::differentiate(math_fn, "y");
+  for (double x : {1.3, -1.3}) {
+    std::function<double(double)> gx = [x](double t) {
+      return math_fn(t, 0.7);
+    };
+    std::function<double(double)> gy = [x](double t) { return math_fn(x, t); };
+    EXPECT_NEAR(fn_x.execute(x, 0.7), finite_difference_tangent(gx, x, eps),
+                tol);
+    EXPECT_NEAR(fn_y.execute(x, 0.7), finite_difference_tangent(gy, 0.7, eps),
+                tol);
+  }
+}
+
+TEST(KokkosMath, Reverse) {
+  const double eps = 1e-6, tol = 1e-5;
+  auto fn_grad = clad::gradient(math_fn);
+  for (double x : {1.3, -1.3}) {
+    double dx = 0, dy = 0;
+    fn_grad.execute(x, 0.7, &dx, &dy);
+    std::function<double(double)> gx = [x](double t) {
+      return math_fn(t, 0.7);
+    };
+    std::function<double(double)> gy = [x](double t) { return math_fn(x, t); };
+    EXPECT_NEAR(dx, finite_difference_tangent(gx, x, eps), tol);
+    EXPECT_NEAR(dy, finite_difference_tangent(gy, 0.7, eps), tol);
+  }
+}