Merge pull request #3139 from stan-dev/normal-quantile-gradients

Improve numerical stability of normal quantile gradients

Author: andrjohns

stan/math/rev/fun/inv_Phi.hpp

@@ -3,8 +3,10 @@
 
 #include <stan/math/rev/meta.hpp>
 #include <stan/math/rev/core.hpp>
-#include <stan/math/prim/fun/constants.hpp>
+#include <stan/math/prim/fun/exp.hpp>
 #include <stan/math/prim/fun/inv_Phi.hpp>
+#include <stan/math/prim/prob/std_normal_lpdf.hpp>
+#include <stan/math/prim/functor/apply_scalar_binary.hpp>
 #include <cmath>
 
 namespace stan {
@@ -19,8 +21,9 @@ namespace math {
  * @return The unit normal inverse cdf evaluated at p
  */
 inline var inv_Phi(const var& p) {
-  return make_callback_var(inv_Phi(p.val()), [p](auto& vi) mutable {
-    p.adj() += vi.adj() * SQRT_TWO_PI / std::exp(-0.5 * vi.val() * vi.val());
+  double val = inv_Phi(p.val());
+  return make_callback_var(val, [p, val](auto& vi) mutable {
+    p.adj() += vi.adj() * exp(-std_normal_lpdf(val));
   });
 }
 
@@ -33,9 +36,11 @@ inline var inv_Phi(const var& p) {
  */
 template <typename T, require_var_matrix_t<T>* = nullptr>
 inline auto inv_Phi(const T& p) {
-  return make_callback_var(inv_Phi(p.val()), [p](auto& vi) mutable {
-    p.adj().array() += vi.adj().array() * SQRT_TWO_PI
-                       / (-0.5 * vi.val().array().square()).exp();
+  auto arena_rtn = to_arena(inv_Phi(p.val()));
+  return make_callback_var(arena_rtn, [p, arena_rtn](auto& vi) mutable {
+    auto deriv
+        = arena_rtn.unaryExpr([](auto x) { return exp(-std_normal_lpdf(x)); });
+    p.adj() += elt_multiply(vi.adj(), deriv);
   });
 }
 

stan/math/rev/prob/std_normal_log_qf.hpp

@@ -3,9 +3,9 @@
 
 #include <stan/math/rev/meta.hpp>
 #include <stan/math/rev/core.hpp>
-#include <stan/math/prim/fun/constants.hpp>
-#include <stan/math/prim/fun/sign.hpp>
 #include <stan/math/prim/prob/std_normal_log_qf.hpp>
+#include <stan/math/prim/functor/apply_scalar_binary.hpp>
+#include <stan/math/prim/fun/elt_multiply.hpp>
 #include <cmath>
 
 namespace stan {
@@ -19,16 +19,18 @@ namespace math {
  */
 template <typename T, require_stan_scalar_or_eigen_t<T>* = nullptr>
 inline auto std_normal_log_qf(const var_value<T>& log_p) {
-  return make_callback_var(
-      std_normal_log_qf(log_p.val()), [log_p](auto& vi) mutable {
-        auto vi_array = as_array_or_scalar(vi.val());
-        auto vi_sign = sign(as_array_or_scalar(vi.adj()));
-
-        const auto& deriv = as_array_or_scalar(log_p).val()
-                            + log(as_array_or_scalar(vi.adj()) * vi_sign)
-                            - NEG_LOG_SQRT_TWO_PI + 0.5 * square(vi_array);
-        as_array_or_scalar(log_p).adj() += vi_sign * exp(deriv);
-      });
+  auto arena_rtn = to_arena(std_normal_log_qf(log_p.val()));
+  return make_callback_var(arena_rtn, [log_p, arena_rtn](auto& vi) mutable {
+    if constexpr (is_eigen<decltype(arena_rtn)>::value) {
+      auto deriv = exp(log_p.val() - arena_rtn.unaryExpr([](auto x) {
+        return std_normal_lpdf(x);
+      }));
+      log_p.adj() += elt_multiply(vi.adj(), deriv);
+    } else {
+      auto deriv = exp(log_p.val() - std_normal_lpdf(arena_rtn));
+      log_p.adj() += vi.adj() * deriv;
+    }
+  });
 }
 
 }  // namespace math