More stable incomplete beta gradient roots

stan/math/fwd/fun/inc_beta.hpp

@@ -3,54 +3,59 @@
 
 #include <stan/math/fwd/core.hpp>
 #include <stan/math/fwd/meta.hpp>
-#include <stan/math/fwd/fun/beta.hpp>
 #include <stan/math/fwd/fun/digamma.hpp>
-#include <stan/math/fwd/fun/exp.hpp>
-#include <stan/math/fwd/fun/lbeta.hpp>
-#include <stan/math/fwd/fun/pow.hpp>
 #include <stan/math/prim/fun/grad_reg_inc_beta.hpp>
 #include <stan/math/prim/fun/inc_beta.hpp>
-#include <cmath>
+#include <stan/math/prim/fun/inc_beta_ddz.hpp>
+#include <stan/math/prim/fun/value_of.hpp>
 
 namespace stan {
 namespace math {
 
-template <typename T>
-inline fvar<T> inc_beta(const fvar<T>& a, const fvar<T>& b, const fvar<T>& x) {
-  T d_a;
-  T d_b;
-  const T beta_ab = beta(a.val_, b.val_);
-  grad_reg_inc_beta(d_a, d_b, a.val_, b.val_, x.val_, digamma(a.val_),
-                    digamma(b.val_), digamma(a.val_ + b.val_), beta_ab);
-  T d_x = pow((1 - x.val_), b.val_ - 1) * pow(x.val_, a.val_ - 1) / beta_ab;
-  return fvar<T>(inc_beta(a.val_, b.val_, x.val_),
-                 a.d_ * d_a + b.d_ * d_b + x.d_ * d_x);
-}
-
-template <typename T>
-inline fvar<T> inc_beta(double a, const fvar<T>& b, const fvar<T>& x) {
-  return inc_beta(fvar<T>(a), b, x);
-}
-template <typename T>
-inline fvar<T> inc_beta(const fvar<T>& a, double b, const fvar<T>& x) {
-  return inc_beta(a, fvar<T>(b), x);
-}
-template <typename T>
-inline fvar<T> inc_beta(const fvar<T>& a, const fvar<T>& b, double x) {
-  return inc_beta(a, b, fvar<T>(x));
-}
-
-template <typename T>
-inline fvar<T> inc_beta(double a, double b, const fvar<T>& x) {
-  return inc_beta(fvar<T>(a), fvar<T>(b), x);
-}
-template <typename T>
-inline fvar<T> inc_beta(const fvar<T>& a, double b, double x) {
-  return inc_beta(a, fvar<T>(b), fvar<T>(x));
-}
-template <typename T>
-inline fvar<T> inc_beta(double a, const fvar<T>& b, double x) {
-  return inc_beta(fvar<T>(a), b, fvar<T>(x));
+/**
+ * The regularized incomplete beta function I_x(a, b) for <code>fvar</code>
+ * arguments. Only the partials of the <code>fvar</code> arguments are
+ * computed. Each argument is an <code>fvar</code> or arithmetic. Mixing with
+ * <code>var</code> is not supported.
+ *
+ * @tparam T1 type of a
+ * @tparam T2 type of b
+ * @tparam T3 type of x
+ * @param a first shape, positive
+ * @param b second shape, positive
+ * @param x upper limit, in [0, 1]
+ * @return I_x(a, b)
+ */
+template <typename T1, typename T2, typename T3,
+          require_all_stan_scalar_t<T1, T2, T3>* = nullptr,
+          require_all_not_var_t<T1, T2, T3>* = nullptr,
+          require_any_fvar_t<T1, T2, T3>* = nullptr>
+inline fvar<partials_return_t<T1, T2, T3>> inc_beta(const T1& a, const T2& b,
+                                                    const T3& x) {
+  using T_return = partials_return_t<T1, T2, T3>;
+  const auto a_val = value_of(a);
+  const auto b_val = value_of(b);
+  const auto x_val = value_of(x);
+  T_return d(0);
+  if constexpr (is_autodiff_v<T1> && is_autodiff_v<T2>) {
+    const auto [d_a, d_b] = grad_reg_inc_beta<GRAD_RIB::BOTH>(
+        a_val, b_val, x_val, digamma(a_val), digamma(b_val),
+        digamma(a_val + b_val), 0.0);
+    d += a.d_ * d_a + b.d_ * d_b;
+  } else if constexpr (is_autodiff_v<T1>) {
+    d += a.d_
+         * grad_reg_inc_beta<GRAD_RIB::G1>(a_val, b_val, x_val, digamma(a_val),
+                                           0.0, digamma(a_val + b_val), 0.0);
+  } else if constexpr (is_autodiff_v<T2>) {
+    d += b.d_
+         * grad_reg_inc_beta<GRAD_RIB::G2>(a_val, b_val, x_val, 0.0,
+                                           digamma(b_val),
+                                           digamma(a_val + b_val), 0.0);
+  }
+  if constexpr (is_autodiff_v<T3>) {
+    d += x.d_ * inc_beta_ddz(a_val, b_val, x_val);
+  }
+  return fvar<T_return>(inc_beta(a_val, b_val, x_val), d);
 }
 
 }  // namespace math

stan/math/prim/fun.hpp

@@ -134,6 +134,9 @@
 #include <stan/math/prim/fun/if_else.hpp>
 #include <stan/math/prim/fun/imag.hpp>
 #include <stan/math/prim/fun/inc_beta.hpp>
+#include <stan/math/prim/fun/inc_beta_dda.hpp>
+#include <stan/math/prim/fun/inc_beta_ddb.hpp>
+#include <stan/math/prim/fun/inc_beta_ddz.hpp>
 #include <stan/math/prim/fun/initialize.hpp>
 #include <stan/math/prim/fun/initialize_fill.hpp>
 #include <stan/math/prim/fun/int_step.hpp>