Accurate lub_constrain derivatives for large positive arguments

stan/math/opencl/rev/constraint/lub_constrain.hpp

@@ -52,8 +52,8 @@ inline var_value<matrix_cl<double>> lub_constrain(T_x&& x, T_lb&& lb,
         auto lb_inf = value_of(lb_arena) == NEGATIVE_INFTY;
         auto ub_inf = value_of(ub_arena) == INFTY;
         auto inv_logit_x = inv_logit(value_of(x_arena));
-        auto one_m_inv_logit_x = 1.0 - inv_logit_x;
         auto exp_x = exp(value_of(x_arena));
+        auto one_m_inv_logit_x = elt_divide(1.0, 1.0 + exp_x);
         auto res_adj_exp_x = elt_multiply(res.adj(), exp_x);
         adjoint_results(x_arena, lb_arena, ub_arena) += expressions(
             select(lb_inf, select(ub_inf, res.adj(), -res_adj_exp_x),
@@ -112,10 +112,10 @@ inline var_value<matrix_cl<double>> lub_constrain(T_x&& x, T_lb&& lb, T_ub&& ub,
         auto lb_inf = value_of(lb_arena) == NEGATIVE_INFTY;
         auto ub_inf = value_of(ub_arena) == INFTY;
         auto inv_logit_x = inv_logit(value_of(x_arena));
-        auto one_m_inv_logit_x = 1.0 - inv_logit_x;
+        auto exp_x = exp(value_of(x_arena));
+        auto one_m_inv_logit_x = elt_divide(1.0, 1.0 + exp_x);
         auto diff = value_of(ub_arena) - value_of(lb_arena);
         auto one_over_diff = elt_divide(1.0, diff);
-        auto exp_x = exp(value_of(x_arena));
         auto res_adj_exp_x = elt_multiply(res.adj(), exp_x);
         adjoint_results(x_arena, lb_arena, ub_arena) += expressions(
             select(lb_inf, select(ub_inf, res.adj(), lp.adj() - res_adj_exp_x),

stan/math/prim/constraint/lub_constrain.hpp

@@ -5,6 +5,7 @@
 #include <stan/math/prim/constraint/lb_constrain.hpp>
 #include <stan/math/prim/constraint/ub_constrain.hpp>
 #include <stan/math/prim/fun/Eigen.hpp>
+#include <stan/math/prim/fun/value_of_rec.hpp>
 #include <stan/math/prim/meta.hpp>
 #include <stan/math/prim/err.hpp>
 #include <stan/math/prim/fun/add.hpp>
@@ -55,7 +56,17 @@ inline auto lub_constrain(T&& x, L&& lb, U&& ub) {
     return ub_constrain(identity_constrain(x, lb), ub);
   } else {
     check_less("lub_constrain", "lb", value_of(lb), value_of(ub));
-    return (ub - lb) * inv_logit(x) + lb;
+    const auto diff = ub - lb;
+    if constexpr (is_autodiff_v<return_type_t<T, L, U>>) {
+      // inv_logit of a non-positive argument only: the derivatives of
+      // 1 - inv_logit(x) cancel for x >> 0
+      if (value_of_rec(x) >= 0) {
+        return ub - diff * inv_logit(-x);
+      }
+      return lb + diff * inv_logit(x);
+    } else {
+      return diff * inv_logit(x) + lb;
+    }
   }
 }
 
@@ -108,8 +119,22 @@ inline auto lub_constrain(T&& x, L&& lb, U&& ub, Lp& lp) {
   } else {
     check_less("lub_constrain", "lb", value_of(lb), value_of(ub));
     const auto diff = ub - lb;
-    lp += add(log(diff), subtract(-abs(x), multiply(2.0, log1p_exp(-abs(x)))));
-    return diff * inv_logit(x) + lb;
+    // a branch instead of abs(x), whose forward-mode derivative is 0 at 0
+    const bool x_nonneg = value_of_rec(x) >= 0;
+    if (x_nonneg) {
+      lp += add(log(diff), subtract(-x, multiply(2.0, log1p_exp(-x))));
+    } else {
+      lp += add(log(diff), subtract(x, multiply(2.0, log1p_exp(x))));
+    }
+    if constexpr (is_autodiff_v<return_type_t<T, L, U>>) {
+      // inv_logit of a non-positive argument only, as above
+      if (x_nonneg) {
+        return ub - diff * inv_logit(-x);
+      }
+      return lb + diff * inv_logit(x);
+    } else {
+      return diff * inv_logit(x) + lb;
+    }
   }
 }