Merge pull request #3033 from nhuurre/sqrt-0-gradient

sqrt(x) gradient when x=0

Author: andrjohns

stan/math/fwd/fun/sqrt.hpp

@@ -16,6 +16,9 @@ namespace math {
 template <typename T>
 inline fvar<T> sqrt(const fvar<T>& x) {
   using std::sqrt;
+  if (value_of_rec(x.val_) == 0.0) {
+    return fvar<T>(sqrt(x.val_), 0.0 * x.d_);
+  }
   return fvar<T>(sqrt(x.val_), 0.5 * x.d_ * inv_sqrt(x.val_));
 }
 

stan/math/rev/fun/sqrt.hpp

@@ -43,7 +43,9 @@ namespace math {
  */
 inline var sqrt(const var& a) {
   return make_callback_var(std::sqrt(a.val()), [a](auto& vi) mutable {
-    a.adj() += vi.adj() / (2.0 * vi.val());
+    if (vi.val() != 0.0) {
+      a.adj() += vi.adj() / (2.0 * vi.val());
+    }
   });
 }
 
@@ -58,7 +60,9 @@ template <typename T, require_var_matrix_t<T>* = nullptr>
 inline auto sqrt(const T& a) {
   return make_callback_var(
       a.val().array().sqrt().matrix(), [a](auto& vi) mutable {
-        a.adj().array() += vi.adj().array() / (2.0 * vi.val_op().array());
+        a.adj().array()
+            += (vi.val_op().array() == 0.0)
+                   .select(0.0, vi.adj().array() / (2.0 * vi.val_op().array()));
       });
 }
 

test/unit/math/mix/fun/distance_test.cpp

@@ -3,6 +3,11 @@
 TEST(MathMixMatFun, distance) {
   auto f
       = [](const auto& x, const auto& y) { return stan::math::distance(x, y); };
+  stan::test::ad_tolerances tols;
+  tols.hessian_hessian_ = 2.0;
+  tols.hessian_fvar_hessian_ = 2.0;
+  tols.grad_hessian_hessian_ = 2.0;
+  tols.grad_hessian_grad_hessian_ = 2.0;
 
   // 0 x 0
   Eigen::VectorXd x0(0);
@@ -13,6 +18,7 @@ TEST(MathMixMatFun, distance) {
   // 1 x 1
   Eigen::VectorXd x1(1);
   x1 << 1;
+  stan::test::expect_ad(tols, f, x1, x1);
   Eigen::VectorXd y1(1);
   y1 << -2.3;
   stan::test::expect_ad(f, x1, y1);
@@ -21,6 +27,7 @@ TEST(MathMixMatFun, distance) {
   // 2 x 2
   Eigen::VectorXd x2(2);
   x2 << 2, -3;
+  stan::test::expect_ad(tols, f, x2, x2);
   Eigen::VectorXd y2(2);
   y2 << -2.3, 1.1;
   stan::test::expect_ad(f, x2, y2);
@@ -29,6 +36,7 @@ TEST(MathMixMatFun, distance) {
   // 3 x 3
   Eigen::VectorXd x(3);
   x << 1, 3, -5;
+  stan::test::expect_ad(tols, f, x, x);
   Eigen::VectorXd y(3);
   y << 4, -2, -1;
   stan::test::expect_ad(f, x, y);

test/unit/math/mix/fun/sqrt_test.cpp

@@ -6,11 +6,21 @@ TEST(mathMixMatFun, sqrt) {
     using stan::math::sqrt;
     return sqrt(x1);
   };
+  auto ff = [](const auto& x1) {
+    using stan::math::sqrt;
+    return sqrt((x1 >= 0.0) ? x1 : -x1);
+  };
   stan::test::expect_common_nonzero_unary_vectorized<
       stan::test::ScalarSupport::Real>(f);
   stan::test::expect_unary_vectorized(f, -6, -5.2, 1.3, 7, 10.7, 36, 1e6);
 
-  // undefined with 0 in denominator
+  stan::test::ad_tolerances tols;
+  tols.hessian_hessian_ = 2.0;
+  tols.hessian_fvar_hessian_ = 2.0;
+  tols.grad_hessian_hessian_ = 2.0;
+  tols.grad_hessian_grad_hessian_ = 2.0;
+  stan::test::expect_ad(tols, ff, 0.0);
+
   stan::test::expect_ad(f, std::complex<double>(0.9, 0.8));
   for (double im : std::vector<double>{-1.3, 2.3}) {
     for (double re : std::vector<double>{-3.6, -0.0, 0.0, 0.5}) {