Improve numerical stability of normal quantile gradients

Author: andrjohns

stan/math/rev/fun/inv_Phi.hpp

@@ -3,8 +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/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 +20,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 +35,12 @@ 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();
+  const auto& arena_rtn = to_arena(inv_Phi(p.val()));
+  return make_callback_var(arena_rtn, [p, arena_rtn](auto& vi) mutable {
+    p.adj() += apply_scalar_binary(vi.adj(), arena_rtn.val(),
+                                  [](const auto& adj, const auto& rtn_val) {
+                                    return adj * exp(-std_normal_lpdf(rtn_val));
+                                  });
   });
 }
 

stan/math/rev/prob/std_normal_log_qf.hpp

@@ -3,9 +3,8 @@
 
 #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_ternary.hpp>
 #include <cmath>
 
 namespace stan {
@@ -19,15 +18,12 @@ 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);
+  const 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 {
+        log_p.adj() += apply_scalar_ternary(
+          [](const auto& adj, const auto& logp_val, const auto& rtn_val) {
+            return adj * exp(logp_val - std_normal_lpdf(rtn_val));
+          }, vi.adj(), log_p.val(), arena_rtn.val());
       });
 }