Add analytical derivative support for std::beta

Author: leetcodez

include/clad/Differentiator/BuiltinDerivatives.h

@@ -1062,6 +1062,61 @@ CUDA_HOST_DEVICE void hypot_pullback(T x, T y, U d_z, T* d_x, T* d_y) {
   *d_y += (y / h) * d_z;
 }
 
+// 7. Special Functions
+#if __cplusplus >= 201703L
+template <typename T> CUDA_HOST_DEVICE inline T clad_beta_primal(T x, T y) {
+#if defined(__cpp_lib_math_special_functions)
+  return ::std::beta(x, y);
+#else
+  return ::std::tgamma(x) * ::std::tgamma(y) / ::std::tgamma(x + y);
+#endif
+}
+template <typename T> CUDA_HOST_DEVICE inline T clad_digamma(T x) {
+  if (x <= 0.0) {
+    if (x == ::std::floor(x))
+      return (T)NAN;
+    return clad_digamma(1.0 - x) -
+           ::std::acos((T)-1.0) / ::std::tan(::std::acos((T)-1.0) * x);
+  }
+  T result = 0.0;
+  while (x < 8.0) {
+    result -= 1.0 / x;
+    x += 1.0;
+  }
+  T inv_x = 1.0 / x;
+  T inv_x2 = inv_x * inv_x;
+  result +=
+      ::std::log(x) - 0.5 * inv_x -
+      inv_x2 * (1.0 / 12.0 -
+                inv_x2 * (1.0 / 120.0 -
+                          inv_x2 * (1.0 / 252.0 - inv_x2 * (1.0 / 240.0))));
+  return result;
+}
+
+template <typename T, typename dT>
+CUDA_HOST_DEVICE ValueAndPushforward<T, dT> beta_pushforward(T x, T y, dT d_x,
+                                                             dT d_y) {
+  T b = clad_beta_primal(x, y);
+  T psi_xy = clad_digamma(x + y);
+  dT pushforward = 0;
+  if (d_x)
+    pushforward += b * (clad_digamma(x) - psi_xy) * d_x;
+  if (d_y)
+    pushforward += b * (clad_digamma(y) - psi_xy) * d_y;
+  return {b, pushforward};
+}
+
+template <typename T, typename U>
+CUDA_HOST_DEVICE void beta_pullback(T x, T y, U d_z, T* d_x, T* d_y) {
+  T b = clad_beta_primal(x, y);
+  T psi_xy = clad_digamma(x + y);
+  if (d_x)
+    *d_x += b * (clad_digamma(x) - psi_xy) * d_z;
+  if (d_y)
+    *d_y += b * (clad_digamma(y) - psi_xy) * d_z;
+}
+#endif
+
 } // namespace std
 
 CUDA_HOST_DEVICE inline ValueAndPushforward<float, float>
@@ -1327,6 +1382,12 @@ using std::pow_pullback;
 using std::pow_pushforward;
 using std::sqrt_pushforward;
 
+// 7. Special Functions
+#if __cplusplus >= 201703L
+using std::beta_pullback;
+using std::beta_pushforward;
+#endif
+
 namespace class_functions {
 template <typename T, typename U>
 void constructor_pullback(ValueAndPushforward<T, U> rhs,

test/Features/stl-cmath.cpp

@@ -123,6 +123,14 @@
 #include <cmath>
 #include <iostream>
 #include <iomanip>
+#if !defined(__cpp_lib_math_special_functions)
+namespace std {
+  template <typename T>
+  inline T beta(T x, T y) {
+    return std::tgamma(x) * std::tgamma(y) / std::tgamma(x + y);
+  }
+}
+#endif
 
 template <typename T>
 T get_tolerance() {
@@ -295,6 +303,10 @@ DEFINE_FUNCTIONS(atanh) // x in [-1,1]
 //
 DEFINE_FUNCTIONS(erf)  // x in (-inf,+inf)
 
+template<typename T> T f_beta(T x){ return std::beta(x,(T)2.0); } // x in (0, +inf)
+inline float f_betaf(float x){ return std::beta(x, 2.0f); }
+inline long double f_betal(long double x){ return std::beta(x, 2.0L); }
+
 int main() {
   // Absolute value
   CHECK(abs);
@@ -352,6 +364,7 @@ int main() {
 
   // Error / Gamma functions
   CHECK_ALL(erf);
+  CHECK_ALL_RANGE(beta, {0.1, 0.5, 1.0, 1.5, 2.0, 3.0, 4.0});
 
   return 0;
 }

test/FirstDerivative/BuiltinDerivatives.C

@@ -6,6 +6,17 @@
 
 #include "clad/Differentiator/Differentiator.h"
 #include "../TestUtils.h"
+#include <cmath>
+
+#if !defined(__cpp_lib_math_special_functions)
+namespace std {
+  // Mock std::beta on Apple platforms so Clad's AST has a target to differentiate
+  inline double beta(double x, double y) {
+    return std::tgamma(x) * std::tgamma(y) / std::tgamma(x + y);
+  }
+}
+#endif
+
 extern "C" int printf(const char* fmt, ...);
 
 
@@ -515,6 +526,14 @@ double f_custom_min(double x, double y) { return std::min(x, y, std::greater<dou
 // CHECK-NEXT:      return _t0.pushforward;
 // CHECK-NEXT:  }
 
+double f_beta(double x, double y) { return std::beta(x, y); }
+// CHECK: double f_beta_darg0(double x, double y) {
+// CHECK-NEXT:     double _d_x = 1;
+// CHECK-NEXT:     double _d_y = 0;
+// CHECK-NEXT:     {{.*}}ValueAndPushforward<double, double> _t0 = {{.*}}beta_pushforward(x, y, _d_x, _d_y);
+// CHECK-NEXT:     return _t0.pushforward;
+// CHECK-NEXT: }
+
 int main () { //expected-no-diagnostics
   float f_result[2];
   double d_result[2];
@@ -695,5 +714,8 @@ int main () { //expected-no-diagnostics
   auto d_custom_min =  clad::differentiate(f_custom_min, 0);
   printf("Result is = %.6f\n", d_custom_min.execute(2, 3)); // CHECK-EXEC: Result is = 0.000000
 
+  auto d_beta = clad::differentiate(f_beta, 0);
+  printf("Result is = %.6f\n", d_beta.execute(2.0, 3.0)); // CHECK-EXEC: Result is = -0.090278
+
   return 0;
 }