diff --git a/stan/math/opencl/rev/constraint/lub_constrain.hpp b/stan/math/opencl/rev/constraint/lub_constrain.hpp index 9fae858693d..08c016ddf31 100644 --- a/stan/math/opencl/rev/constraint/lub_constrain.hpp +++ b/stan/math/opencl/rev/constraint/lub_constrain.hpp @@ -52,8 +52,8 @@ inline var_value> 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> 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), diff --git a/stan/math/prim/constraint/lub_constrain.hpp b/stan/math/prim/constraint/lub_constrain.hpp index e7815664e46..340f7020b87 100644 --- a/stan/math/prim/constraint/lub_constrain.hpp +++ b/stan/math/prim/constraint/lub_constrain.hpp @@ -17,6 +17,7 @@ #include #include #include +#include #include namespace stan { @@ -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>) { + // 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>) { + // 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; + } } } diff --git a/stan/math/rev/constraint/lub_constrain.hpp b/stan/math/rev/constraint/lub_constrain.hpp index df5d9f9dd46..84850876351 100644 --- a/stan/math/rev/constraint/lub_constrain.hpp +++ b/stan/math/rev/constraint/lub_constrain.hpp @@ -3,6 +3,8 @@ #include #include +#include +#include #include #include #include @@ -10,6 +12,25 @@ namespace stan { namespace math { +namespace internal { +/** + * Return in arena memory the inverse logit of x given exp(-|x|). + * + * The reverse passes compute inv_logit(-x) from the result without a + * division: exp(-|x|) * inv_logit(x) where inv_logit(x) >= 1/2 (x >= 0), + * and 1 - inv_logit(x) otherwise, which does not cancel there. + * + * @tparam T_x type of the array of x + * @tparam T_e type of the array of exp(-|x|) + * @param x x + * @param exp_neg_abs_x exp(-|x|) + * @return inverse logit of x + */ +template +inline auto lub_inv_logit(const T_x& x, const T_e& exp_neg_abs_x) { + return to_arena((x >= 0).select(1.0, exp_neg_abs_x) / (1.0 + exp_neg_abs_x)); +} +} // namespace internal /** * Return the lower and upper-bounded scalar derived by @@ -48,20 +69,24 @@ inline auto lub_constrain(const T& x, const L& lb, const U& ub) { } else { check_less("lub_constrain", "lb", lb_val, ub_val); auto diff = ub_val - lb_val; - double inv_logit_x = inv_logit(value_of(x)); - return make_callback_var(diff * inv_logit_x + lb_val, - [x, ub, lb, diff, inv_logit_x](auto& vi) mutable { - if constexpr (is_autodiff_v) { - x.adj() += vi.adj() * diff * inv_logit_x - * (1.0 - inv_logit_x); - } - if constexpr (is_autodiff_v) { - lb.adj() += vi.adj() * (1.0 - inv_logit_x); - } - if constexpr (is_autodiff_v) { - ub.adj() += vi.adj() * inv_logit_x; - } - }); + const double x_val = value_of(x); + const double e = std::exp(-std::fabs(x_val)); + const double inv_logit_x = (x_val >= 0 ? 1.0 : e) / (1.0 + e); + const double inv_logit_neg_x + = x_val >= 0 ? e * inv_logit_x : 1.0 - inv_logit_x; + return make_callback_var( + diff * inv_logit_x + lb_val, + [x, ub, lb, diff, inv_logit_x, inv_logit_neg_x](auto& vi) mutable { + if constexpr (is_autodiff_v) { + x.adj() += vi.adj() * diff * inv_logit_x * inv_logit_neg_x; + } + if constexpr (is_autodiff_v) { + lb.adj() += vi.adj() * inv_logit_neg_x; + } + if constexpr (is_autodiff_v) { + ub.adj() += vi.adj() * inv_logit_x; + } + }); } } @@ -116,25 +141,26 @@ inline auto lub_constrain(const T& x, const L& lb, const U& ub, return ub_constrain(identity_constrain(x, lb), ub, lp); } else { check_less("lub_constrain", "lb", lb_val, ub_val); - auto neg_abs_x = -abs(value_of(x)); auto diff = ub_val - lb_val; - double inv_logit_x = inv_logit(value_of(x)); - lp += (log(diff) + (neg_abs_x - (2.0 * log1p_exp(neg_abs_x)))); + const double x_val = value_of(x); + const double e = std::exp(-std::fabs(x_val)); + const double inv_logit_x = (x_val >= 0 ? 1.0 : e) / (1.0 + e); + const double inv_logit_neg_x + = x_val >= 0 ? e * inv_logit_x : 1.0 - inv_logit_x; + lp += (log(diff) + (-std::fabs(x_val) - (2.0 * log1p(e)))); return make_callback_var( diff * inv_logit_x + lb_val, - [x, ub, lb, diff, lp, inv_logit_x](auto& vi) mutable { + [x, ub, lb, diff, lp, inv_logit_x, inv_logit_neg_x](auto& vi) mutable { if constexpr (is_autodiff_v) { - x.adj() += vi.adj() * diff * inv_logit_x * (1.0 - inv_logit_x) + x.adj() += vi.adj() * diff * inv_logit_x * inv_logit_neg_x + lp.adj() * (1.0 - 2.0 * inv_logit_x); } if constexpr (is_autodiff_v && is_autodiff_v) { const auto one_over_diff = 1.0 / diff; - lb.adj() - += vi.adj() * (1.0 - inv_logit_x) + -one_over_diff * lp.adj(); + lb.adj() += vi.adj() * inv_logit_neg_x + -one_over_diff * lp.adj(); ub.adj() += vi.adj() * inv_logit_x + one_over_diff * lp.adj(); } else if constexpr (is_autodiff_v) { - lb.adj() - += vi.adj() * (1.0 - inv_logit_x) + (-1.0 / diff) * lp.adj(); + lb.adj() += vi.adj() * inv_logit_neg_x + (-1.0 / diff) * lp.adj(); } else if constexpr (is_autodiff_v) { ub.adj() += vi.adj() * inv_logit_x + (1.0 / diff) * lp.adj(); } @@ -165,20 +191,26 @@ inline auto lub_constrain(const T& x, const L& lb, const U& ub) { arena_t arena_x = x; check_less("lub_constrain", "lb", lb_val, ub_val); const auto diff = ub_val - lb_val; - auto inv_logit_x = to_arena(inv_logit(arena_x.val().array())); + auto exp_neg_abs_x = to_arena((-(arena_x.val().array()).abs()).exp()); + auto inv_logit_x + = internal::lub_inv_logit(arena_x.val().array(), exp_neg_abs_x); arena_t ret = diff * inv_logit_x + lb_val; - reverse_pass_callback([arena_x, ub, lb, ret, diff, inv_logit_x]() mutable { - if constexpr (is_autodiff_v) { - arena_x.adj().array() - += ret.adj().array() * diff * inv_logit_x * (1.0 - inv_logit_x); - } - if constexpr (is_autodiff_v) { - lb.adj() += (ret.adj().array() * (1.0 - inv_logit_x)).sum(); - } - if constexpr (is_autodiff_v) { - ub.adj() += (ret.adj().array() * inv_logit_x).sum(); - } - }); + reverse_pass_callback( + [arena_x, ub, lb, ret, diff, exp_neg_abs_x, inv_logit_x]() mutable { + const auto inv_logit_neg_x + = (inv_logit_x >= 0.5) + .select(exp_neg_abs_x * inv_logit_x, 1.0 - inv_logit_x); + if constexpr (is_autodiff_v) { + arena_x.adj().array() + += ret.adj().array() * diff * inv_logit_x * inv_logit_neg_x; + } + if constexpr (is_autodiff_v) { + lb.adj() += (ret.adj().array() * inv_logit_neg_x).sum(); + } + if constexpr (is_autodiff_v) { + ub.adj() += (ret.adj().array() * inv_logit_x).sum(); + } + }); return ret_type(ret); } } @@ -206,27 +238,33 @@ inline auto lub_constrain(const T& x, const L& lb, const U& ub, } else { check_less("lub_constrain", "lb", lb_val, ub_val); arena_t arena_x = x; - auto neg_abs_x = to_arena(-(value_of(arena_x).array()).abs()); auto diff = ub_val - lb_val; - lp += (log(diff) + (neg_abs_x - (2.0 * log1p_exp(neg_abs_x)))).sum(); - auto inv_logit_x = to_arena(inv_logit(value_of(arena_x).array())); + auto exp_neg_abs_x = to_arena((-(value_of(arena_x).array()).abs()).exp()); + auto inv_logit_x + = internal::lub_inv_logit(value_of(arena_x).array(), exp_neg_abs_x); + // separate sums: log1p vectorizes only without the AoS view of x + lp += (log(diff) - 2.0 * exp_neg_abs_x.log1p()).sum() + - value_of(arena_x).array().abs().sum(); arena_t ret = diff * inv_logit_x + lb_val; reverse_pass_callback( - [arena_x, ub, lb, ret, lp, diff, inv_logit_x]() mutable { + [arena_x, ub, lb, ret, lp, diff, exp_neg_abs_x, inv_logit_x]() mutable { + const auto inv_logit_neg_x + = (inv_logit_x >= 0.5) + .select(exp_neg_abs_x * inv_logit_x, 1.0 - inv_logit_x); if constexpr (is_autodiff_v) { arena_x.adj().array() - += ret.adj().array() * diff * inv_logit_x * (1.0 - inv_logit_x) + += ret.adj().array() * diff * inv_logit_x * inv_logit_neg_x + lp.adj() * (1.0 - 2.0 * inv_logit_x); } if constexpr (is_autodiff_v && is_autodiff_v) { const auto lp_calc = lp.adj() * ret.size(); const auto one_over_diff = 1.0 / diff; - lb.adj() += (ret.adj().array() * (1.0 - inv_logit_x)).sum() + lb.adj() += (ret.adj().array() * inv_logit_neg_x).sum() + -one_over_diff * lp_calc; ub.adj() += (ret.adj().array() * inv_logit_x).sum() + one_over_diff * lp_calc; } else if constexpr (is_autodiff_v) { - lb.adj() += (ret.adj().array() * (1.0 - inv_logit_x)).sum() + lb.adj() += (ret.adj().array() * inv_logit_neg_x).sum() + -(1.0 / diff) * lp.adj() * ret.size(); } else if constexpr (is_autodiff_v) { ub.adj() += (ret.adj().array() * inv_logit_x).sum() @@ -259,15 +297,20 @@ inline auto lub_constrain(const T& x, const L& lb, const U& ub) { check_less("lub_constrain", "lb", lb_val, ub_val); auto is_lb_inf = to_arena((lb_val == NEGATIVE_INFTY)); auto diff = to_arena(ub_val - lb_val); - auto inv_logit_x = to_arena(inv_logit(value_of(arena_x).array())); + auto exp_neg_abs_x = to_arena((-(value_of(arena_x).array()).abs()).exp()); + auto inv_logit_x + = internal::lub_inv_logit(value_of(arena_x).array(), exp_neg_abs_x); arena_t ret = (is_lb_inf).select( ub_val - value_of(arena_x).array().exp(), diff * inv_logit_x + lb_val); - reverse_pass_callback([arena_x, ub, arena_lb, ret, diff, inv_logit_x, - is_lb_inf]() mutable { + reverse_pass_callback([arena_x, ub, arena_lb, ret, diff, exp_neg_abs_x, + inv_logit_x, is_lb_inf]() mutable { + const auto inv_logit_neg_x + = (inv_logit_x >= 0.5) + .select(exp_neg_abs_x * inv_logit_x, 1.0 - inv_logit_x); if constexpr (is_autodiff_v) { arena_x.adj().array() += (is_lb_inf).select( ret.adj().array() * -value_of(arena_x).array().exp(), - ret.adj().array() * diff * inv_logit_x * (1.0 - inv_logit_x)); + ret.adj().array() * diff * inv_logit_x * inv_logit_neg_x); } if constexpr (is_autodiff_v) { ub.adj() @@ -277,7 +320,7 @@ inline auto lub_constrain(const T& x, const L& lb, const U& ub) { } if constexpr (is_autodiff_v) { arena_lb.adj().array() - += (is_lb_inf).select(0, ret.adj().array() * (1.0 - inv_logit_x)); + += (is_lb_inf).select(0, ret.adj().array() * inv_logit_neg_x); } }); return ret_type(ret); @@ -308,28 +351,30 @@ inline auto lub_constrain(const T& x, const L& lb, const U& ub, check_less("lub_constrain", "lb", lb_val, ub_val); auto is_lb_inf = to_arena((lb_val == NEGATIVE_INFTY)); auto diff = to_arena(ub_val - lb_val); - auto neg_abs_x = to_arena(-arena_x_val.abs()); - auto inv_logit_x = to_arena(inv_logit(arena_x_val)); + auto exp_neg_abs_x = to_arena((-(arena_x_val).abs()).exp()); + auto inv_logit_x = internal::lub_inv_logit(arena_x_val, exp_neg_abs_x); arena_t ret = (is_lb_inf).select(ub_val - arena_x_val.exp(), diff * inv_logit_x + lb_val); lp += (is_lb_inf) - .select(arena_x_val, - log(diff) + (neg_abs_x - (2.0 * log1p_exp(neg_abs_x)))) + .select(arena_x_val, log(diff) + + (-arena_x_val.abs() + - (2.0 * exp_neg_abs_x.log1p()))) .sum(); reverse_pass_callback([arena_x, arena_x_val, ub, arena_lb, ret, lp, diff, - inv_logit_x, is_lb_inf]() mutable { + exp_neg_abs_x, inv_logit_x, is_lb_inf]() mutable { + const auto inv_logit_neg_x + = (inv_logit_x >= 0.5) + .select(exp_neg_abs_x * inv_logit_x, 1.0 - inv_logit_x); const auto lp_adj = lp.adj(); if constexpr (is_autodiff_v) { - const auto x_sign = arena_x_val.sign().eval(); arena_x.adj().array() += (is_lb_inf).select( ret.adj().array() * -arena_x_val.exp() + lp_adj, - ret.adj().array() * diff * inv_logit_x * (1.0 - inv_logit_x) + ret.adj().array() * diff * inv_logit_x * inv_logit_neg_x + lp.adj() * (1.0 - 2.0 * inv_logit_x)); } if constexpr (is_autodiff_v) { - arena_lb.adj().array() - += (is_lb_inf).select(0, ret.adj().array() * (1.0 - inv_logit_x) - + -(1.0 / diff) * lp_adj); + arena_lb.adj().array() += (is_lb_inf).select( + 0, ret.adj().array() * inv_logit_neg_x + -(1.0 / diff) * lp_adj); } if constexpr (is_autodiff_v) { ub.adj() @@ -366,20 +411,25 @@ inline auto lub_constrain(const T& x, const L& lb, const U& ub) { check_less("lub_constrain", "lb", lb_val, ub_val); auto is_ub_inf = to_arena((ub_val == INFTY)); auto diff = to_arena(ub_val - lb_val); - auto inv_logit_x = to_arena(inv_logit(arena_x_val.array())); + auto exp_neg_abs_x = to_arena((-(arena_x_val.array()).abs()).exp()); + auto inv_logit_x + = internal::lub_inv_logit(arena_x_val.array(), exp_neg_abs_x); arena_t ret = (is_ub_inf).select( arena_x_val.array().exp() + lb_val, diff * inv_logit_x + lb_val); reverse_pass_callback([arena_x, arena_x_val, arena_ub, lb, ret, is_ub_inf, - inv_logit_x, diff]() mutable { + exp_neg_abs_x, inv_logit_x, diff]() mutable { + const auto inv_logit_neg_x + = (inv_logit_x >= 0.5) + .select(exp_neg_abs_x * inv_logit_x, 1.0 - inv_logit_x); if constexpr (is_autodiff_v) { arena_x.adj().array() += (is_ub_inf).select( ret.adj().array() * arena_x_val.array().exp(), - ret.adj().array() * diff * inv_logit_x * (1.0 - inv_logit_x)); + ret.adj().array() * diff * inv_logit_x * inv_logit_neg_x); } if constexpr (is_autodiff_v) { lb.adj() += (is_ub_inf) .select(ret.adj().array(), - ret.adj().array() * (1.0 - inv_logit_x)) + ret.adj().array() * inv_logit_neg_x) .sum(); } if constexpr (is_autodiff_v) { @@ -415,27 +465,34 @@ inline auto lub_constrain(const T& x, const L& lb, const U& ub, check_less("lub_constrain", "lb", lb_val, ub_val); auto is_ub_inf = to_arena((ub_val.array() == INFTY)); auto diff = to_arena(ub_val.array() - lb_val); - auto neg_abs_x = to_arena(-(arena_x_val.array()).abs()); + auto exp_neg_abs_x = to_arena((-(arena_x_val.array()).abs()).exp()); + auto inv_logit_x + = internal::lub_inv_logit(arena_x_val.array(), exp_neg_abs_x); lp += (is_ub_inf) .select(arena_x_val.array(), - log(diff) + (neg_abs_x - (2.0 * log1p_exp(neg_abs_x)))) + log(diff) + + (-(arena_x_val.array()).abs() + - (2.0 * exp_neg_abs_x.log1p()))) .sum(); - auto inv_logit_x = to_arena(inv_logit(arena_x_val.array())); arena_t ret = (is_ub_inf).select( arena_x_val.array().exp() + lb_val, diff * inv_logit_x + lb_val); - reverse_pass_callback([arena_x, arena_x_val, diff, inv_logit_x, arena_ub, - lb, ret, lp, is_ub_inf]() mutable { + reverse_pass_callback([arena_x, arena_x_val, diff, exp_neg_abs_x, + inv_logit_x, arena_ub, lb, ret, lp, + is_ub_inf]() mutable { + const auto inv_logit_neg_x + = (inv_logit_x >= 0.5) + .select(exp_neg_abs_x * inv_logit_x, 1.0 - inv_logit_x); const auto lp_adj = lp.adj(); if constexpr (is_autodiff_v) { arena_x.adj().array() += (is_ub_inf).select( ret.adj().array() * arena_x_val.array().exp() + lp_adj, - ret.adj().array() * diff * inv_logit_x * (1.0 - inv_logit_x) + ret.adj().array() * diff * inv_logit_x * inv_logit_neg_x + lp.adj() * (1.0 - 2.0 * inv_logit_x)); } if constexpr (is_autodiff_v) { lb.adj() += (is_ub_inf) .select(ret.adj().array(), - ret.adj().array() * (1.0 - inv_logit_x) + ret.adj().array() * inv_logit_neg_x + -(1.0 / diff) * lp_adj) .sum(); } @@ -464,7 +521,9 @@ inline auto lub_constrain(const T& x, const L& lb, const U& ub) { auto lb_val = value_of(arena_lb).array(); auto ub_val = value_of(arena_ub).array(); check_less("lub_constrain", "lb", lb_val, ub_val); - auto inv_logit_x = to_arena(inv_logit(arena_x_val.array())); + auto exp_neg_abs_x = to_arena((-(arena_x_val.array()).abs()).exp()); + auto inv_logit_x + = internal::lub_inv_logit(arena_x_val.array(), exp_neg_abs_x); auto is_lb_inf = to_arena((lb_val == NEGATIVE_INFTY)); auto is_ub_inf = to_arena((ub_val == INFTY)); auto is_lb_ub_inf = to_arena(is_lb_inf && is_ub_inf); @@ -477,18 +536,21 @@ inline auto lub_constrain(const T& x, const L& lb, const U& ub) { ub_val - arena_x.val().array().exp(), (is_ub_inf).select(arena_x_val.array().exp() + lb_val, diff * inv_logit_x + lb_val))); - reverse_pass_callback([arena_x, arena_x_val, inv_logit_x, arena_ub, arena_lb, - diff, ret, is_ub_inf, is_lb_inf, + reverse_pass_callback([arena_x, arena_x_val, exp_neg_abs_x, inv_logit_x, + arena_ub, arena_lb, diff, ret, is_ub_inf, is_lb_inf, is_lb_ub_inf]() mutable { + const auto inv_logit_neg_x + = (inv_logit_x >= 0.5) + .select(exp_neg_abs_x * inv_logit_x, 1.0 - inv_logit_x); // The most likely case is neither of them are infinity const bool is_none_inf = !(is_lb_inf.any() || is_ub_inf.any()); if (is_none_inf) { if constexpr (is_autodiff_v) { arena_x.adj().array() - += ret.adj().array() * diff * inv_logit_x * (1.0 - inv_logit_x); + += ret.adj().array() * diff * inv_logit_x * inv_logit_neg_x; } if constexpr (is_autodiff_v) { - arena_lb.adj().array() += ret.adj().array() * (1.0 - inv_logit_x); + arena_lb.adj().array() += ret.adj().array() * inv_logit_neg_x; } if constexpr (is_autodiff_v) { arena_ub.adj().array() += ret.adj().array() * inv_logit_x; @@ -504,7 +566,7 @@ inline auto lub_constrain(const T& x, const L& lb, const U& ub) { (is_ub_inf).select( ret.adj().array() * arena_x_val.array().exp(), ret.adj().array() * diff * inv_logit_x - * (1.0 - inv_logit_x)))); + * inv_logit_neg_x))); } if constexpr (is_autodiff_v) { arena_ub.adj().array() += (is_ub_inf).select( @@ -514,7 +576,7 @@ inline auto lub_constrain(const T& x, const L& lb, const U& ub) { if constexpr (is_autodiff_v) { arena_lb.adj().array() += (is_lb_inf).select( 0, (is_ub_inf).select(ret.adj().array(), - ret.adj().array() * (1.0 - inv_logit_x))); + ret.adj().array() * inv_logit_neg_x)); } } }); @@ -538,7 +600,9 @@ inline auto lub_constrain(const T& x, const L& lb, const U& ub, auto lb_val = value_of(arena_lb).array(); auto ub_val = value_of(arena_ub).array(); check_less("lub_constrain", "lb", lb_val, ub_val); - auto inv_logit_x = to_arena(inv_logit(arena_x_val.array())); + auto exp_neg_abs_x = to_arena((-(arena_x_val.array()).abs()).exp()); + auto inv_logit_x + = internal::lub_inv_logit(arena_x_val.array(), exp_neg_abs_x); auto is_lb_inf = to_arena((lb_val == NEGATIVE_INFTY)); auto is_ub_inf = to_arena((ub_val == INFTY)); auto is_lb_ub_inf = to_arena(is_lb_inf && is_ub_inf); @@ -551,30 +615,31 @@ inline auto lub_constrain(const T& x, const L& lb, const U& ub, ub_val - arena_x.val().array().exp(), (is_ub_inf).select(arena_x_val.array().exp() + lb_val, diff * inv_logit_x + lb_val))); - auto neg_abs_x = to_arena(-(arena_x_val.array()).abs()); lp += (is_lb_ub_inf) - .select( - 0, - (is_lb_inf || is_ub_inf) - .select( - arena_x_val.array(), - log(diff) + (neg_abs_x - (2.0 * log1p_exp(neg_abs_x))))) + .select(0, (is_lb_inf || is_ub_inf) + .select(arena_x_val.array(), + log(diff) + + (-(arena_x_val.array()).abs() + - (2.0 * exp_neg_abs_x.log1p())))) .sum(); - reverse_pass_callback([arena_x, arena_x_val, inv_logit_x, arena_ub, arena_lb, - diff, ret, is_ub_inf, is_lb_inf, is_lb_ub_inf, - lp]() mutable { + reverse_pass_callback([arena_x, arena_x_val, exp_neg_abs_x, inv_logit_x, + arena_ub, arena_lb, diff, ret, is_ub_inf, is_lb_inf, + is_lb_ub_inf, lp]() mutable { + const auto inv_logit_neg_x + = (inv_logit_x >= 0.5) + .select(exp_neg_abs_x * inv_logit_x, 1.0 - inv_logit_x); const auto lp_adj = lp.adj(); // The most likely case is neither of them are infinity const bool is_none_inf = !(is_lb_inf.any() || is_ub_inf.any()); if (is_none_inf) { if constexpr (is_autodiff_v) { arena_x.adj().array() - += ret.adj().array() * diff * inv_logit_x * (1.0 - inv_logit_x) + += ret.adj().array() * diff * inv_logit_x * inv_logit_neg_x + lp.adj() * (1.0 - 2.0 * inv_logit_x); } if constexpr (is_autodiff_v) { arena_lb.adj().array() - += ret.adj().array() * (1.0 - inv_logit_x) + -(1.0 / diff) * lp_adj; + += ret.adj().array() * inv_logit_neg_x + -(1.0 / diff) * lp_adj; } if constexpr (is_autodiff_v) { arena_ub.adj().array() @@ -593,13 +658,13 @@ inline auto lub_constrain(const T& x, const L& lb, const U& ub, ret.adj().array() * arena_x_val.array().exp() + lp_adj, ret.adj().array() * diff * inv_logit_x - * (1.0 - inv_logit_x) + * inv_logit_neg_x + lp.adj() * (1.0 - 2.0 * inv_logit_x)))); } if constexpr (is_autodiff_v) { arena_lb.adj().array() += (is_lb_inf).select( 0, (is_ub_inf).select(ret.adj().array(), - ret.adj().array() * (1.0 - inv_logit_x) + ret.adj().array() * inv_logit_neg_x + -(1.0 / diff) * lp_adj)); } if constexpr (is_autodiff_v) { diff --git a/test/unit/math/mix/constraint/lub_constrain_fwd_test.cpp b/test/unit/math/mix/constraint/lub_constrain_fwd_test.cpp new file mode 100644 index 00000000000..079b7de11db --- /dev/null +++ b/test/unit/math/mix/constraint/lub_constrain_fwd_test.cpp @@ -0,0 +1,117 @@ +#include +#include +#include +#include +#include + +namespace lub_constrain_fwd_test { + +// For lb = -2, ub = 3: x, lub_constrain(x, lb, ub), its derivatives with +// respect to lb, ub and x, its second derivative with respect to x, and the +// first and second derivatives of the log Jacobian with respect to x. +// mpmath, 50 digits, rounded to double. +struct ref { + double x; + double val; + double d_lb; + double d_ub; + double d_x; + double d_xx; + double lp_x; + double lp_xx; +}; + +// clang-format off +// NOLINTBEGIN(whitespace/line_length) +inline std::vector refs() { + return { + {0, 0.5, 0.5, 0.5, 1.25, 0, 0, -0.5}, + {1, 1.6552928931500244, 0.2689414213699951, 0.7310585786300049, 0.9830596662074093, -0.45428873836474204, -0.46211715726000974, -0.3932238664829637}, + {-1, -0.6552928931500244, 0.7310585786300049, 0.2689414213699951, 0.9830596662074093, 0.45428873836474204, 0.46211715726000974, -0.3932238664829637}, + {20, 2.9999999896942318, 2.0611536181902037e-09, 0.9999999979388464, 1.0305768069709247e-08, -1.0305768027225704e-08, -0.9999999958776927, -4.122307227883698e-09}, + {-20, -1.999999989694232, 0.9999999979388464, 2.0611536181902037e-09, 1.0305768069709247e-08, 1.0305768027225704e-08, 0.9999999958776927, -4.122307227883698e-09}, + {37, 2.9999999999999996, 8.533047625744065e-17, 0.9999999999999999, 4.2665238128720323e-16, -4.2665238128720313e-16, -0.9999999999999998, -1.706609525148813e-16}, + {-37, -1.9999999999999996, 0.9999999999999999, 8.533047625744065e-17, 4.2665238128720323e-16, 4.2665238128720313e-16, 0.9999999999999998, -1.706609525148813e-16}, + {40, 3.0, 4.248354255291589e-18, 1.0, 2.1241771276457944e-17, -2.1241771276457944e-17, -1.0, -8.496708510583178e-18}, + {-40, -2.0, 1.0, 4.248354255291589e-18, 2.1241771276457944e-17, 2.1241771276457944e-17, 1.0, -8.496708510583178e-18}, + {100, 3.0, 3.720075976020836e-44, 1.0, 1.860037988010418e-43, -1.860037988010418e-43, -1.0, -7.440151952041672e-44}, + {-100, -2.0, 1.0, 3.720075976020836e-44, 1.860037988010418e-43, 1.860037988010418e-43, 1.0, -7.440151952041672e-44}, + {700, 3.0, 9.85967654375977e-305, 1.0, 4.9298382718798856e-304, -4.9298382718798856e-304, -1.0, -1.971935308751954e-304}, + {-700, -2.0, 1.0, 9.85967654375977e-305, 4.9298382718798856e-304, 4.9298382718798856e-304, 1.0, -1.971935308751954e-304}}; +} +// NOLINTEND +// clang-format on + +// Relative check, absolute below the smallest normal double. +inline void expect_rel(double expected, double actual, double x, + const char* what) { + EXPECT_NEAR( + expected, actual, + 1e-13 + * std::fmax(std::fabs(expected), std::numeric_limits::min())) + << what << " at x = " << x; +} + +} // namespace lub_constrain_fwd_test + +TEST(mathMixConstraint, lub_constrain_fwd_derivatives_tails) { + using lub_constrain_fwd_test::expect_rel; + using stan::math::fvar; + using stan::math::lub_constrain; + using stan::math::var; + const double lb = -2.0; + const double ub = 3.0; + for (const auto& r : lub_constrain_fwd_test::refs()) { + expect_rel(r.val, lub_constrain(r.x, lb, ub), r.x, "value"); + expect_rel(r.d_x, lub_constrain(fvar(r.x, 1), lb, ub).d_, r.x, + "d/dx"); + expect_rel(r.d_lb, lub_constrain(r.x, fvar(lb, 1), ub).d_, r.x, + "d/dlb"); + expect_rel(r.d_ub, lub_constrain(r.x, lb, fvar(ub, 1)).d_, r.x, + "d/dub"); + + fvar> x_ff(fvar(r.x, 1), fvar(1, 0)); + fvar> y_ff = lub_constrain(x_ff, lb, ub); + expect_rel(r.d_x, y_ff.d_.val_, r.x, "d/dx, fvar>"); + expect_rel(r.d_xx, y_ff.d_.d_, r.x, "d2/dx2, fvar>"); + + fvar x_fv(r.x, 1); + fvar y_fv = lub_constrain(x_fv, lb, ub); + y_fv.d_.grad(); + expect_rel(r.d_xx, x_fv.val_.adj(), r.x, "d2/dx2, fvar"); + stan::math::recover_memory(); + } +} + +TEST(mathMixConstraint, lub_constrain_fwd_derivatives_tails_lp) { + using lub_constrain_fwd_test::expect_rel; + using stan::math::fvar; + using stan::math::lub_constrain; + using stan::math::var; + const double lb = -2.0; + const double ub = 3.0; + for (const auto& r : lub_constrain_fwd_test::refs()) { + double lp_d = 0; + expect_rel(r.val, lub_constrain(r.x, lb, ub, lp_d), r.x, "value, lp"); + + fvar lp_lb = 0; + expect_rel(r.d_lb, lub_constrain(r.x, fvar(lb, 1), ub, lp_lb).d_, + r.x, "d/dlb, lp"); + + fvar> x_ff(fvar(r.x, 1), fvar(1, 0)); + fvar> lp_ff = 0; + fvar> y_ff = lub_constrain(x_ff, lb, ub, lp_ff); + expect_rel(r.d_x, y_ff.d_.val_, r.x, "d/dx, lp"); + expect_rel(r.d_xx, y_ff.d_.d_, r.x, "d2/dx2, lp"); + expect_rel(r.lp_x, lp_ff.d_.val_, r.x, "log Jacobian d/dx"); + expect_rel(r.lp_xx, lp_ff.d_.d_, r.x, + "log Jacobian d2/dx2, fvar>"); + + fvar x_fv(r.x, 1); + fvar lp_fv = 0; + lub_constrain(x_fv, lb, ub, lp_fv); + lp_fv.d_.grad(); + expect_rel(r.lp_xx, x_fv.val_.adj(), r.x, "log Jacobian d2/dx2, fvar"); + stan::math::recover_memory(); + } +} diff --git a/test/unit/math/opencl/rev/constraint/lub_constrain_test.cpp b/test/unit/math/opencl/rev/constraint/lub_constrain_test.cpp index 121faaf1e81..686409424c3 100644 --- a/test/unit/math/opencl/rev/constraint/lub_constrain_test.cpp +++ b/test/unit/math/opencl/rev/constraint/lub_constrain_test.cpp @@ -3,6 +3,8 @@ #include #include #include +#include +#include auto lub_constrain_functor = [](const auto& a, const auto& b, const auto& c) { return stan::math::lub_constrain(a, b, c); @@ -125,4 +127,63 @@ TEST(OpenCLLubConstrain, prim_rev_values_large) { b_scal, c_scal); } +TEST(OpenCLLubConstrain, rev_derivatives_tails) { + using stan::math::from_matrix_cl; + using stan::math::matrix_cl; + using stan::math::to_matrix_cl; + using stan::math::var; + using stan::math::var_value; + // mpmath at 50 digits, lb = -2, ub = 3: dy/dx = 5 inv_logit(x) inv_logit(-x), + // dy/dlb = inv_logit(-x), dlogJ/dx = inv_logit(-x) - inv_logit(x) + Eigen::VectorXd x(7); + x << -37, -20, 0, 20, 30, 37, 40; + Eigen::VectorXd dx(7); + dx << 4.2665238128720322e-16, 1.0305768069709247e-8, 1.25, + 1.0305768069709247e-8, 4.6788114844192117e-13, 4.2665238128720322e-16, + 2.1241771276457945e-17; + Eigen::VectorXd dlb(7); + dlb << 9.9999999999999991e-1, 9.9999999793884638e-1, 0.5, + 2.0611536181902036e-9, 9.357622968839299e-14, 8.5330476257440651e-17, + 4.248354255291589e-18; + Eigen::VectorXd dlogj(7); + dlogj << 9.9999999999999983e-1, 9.9999999587769276e-1, 0.0, + -9.9999999587769276e-1, -9.9999999999981285e-1, -9.9999999999999983e-1, + -9.9999999999999999e-1; + Eigen::VectorXd lb = Eigen::VectorXd::Constant(7, -2.0); + Eigen::VectorXd ub = Eigen::VectorXd::Constant(7, 3.0); + // relative, absolute below the smallest normal double + auto tol = [](double expected) { + return 1e-13 + * std::fmax(std::fabs(expected), std::numeric_limits::min()); + }; + for (bool use_lp : {false, true}) { + var_value> x_cl = to_matrix_cl(x); + var_value> lb_cl = to_matrix_cl(lb); + var_value> ub_cl = to_matrix_cl(ub); + var lp = 0; + var s + = use_lp + ? stan::math::sum( + stan::math::lub_constrain(x_cl, lb_cl, ub_cl, lp)) + : stan::math::sum(stan::math::lub_constrain(x_cl, lb_cl, ub_cl)); + s.grad(); + Eigen::VectorXd x_adj = from_matrix_cl(x_cl.adj()); + Eigen::VectorXd lb_adj = from_matrix_cl(lb_cl.adj()); + for (int i = 0; i < x.size(); ++i) { + EXPECT_NEAR(dx(i), x_adj(i), tol(dx(i))) << "dx at x = " << x(i); + EXPECT_NEAR(dlb(i), lb_adj(i), tol(dlb(i))) << "dlb at x = " << x(i); + } + if (use_lp) { + stan::math::set_zero_all_adjoints(); + lp.grad(); + x_adj = from_matrix_cl(x_cl.adj()); + for (int i = 0; i < x.size(); ++i) { + EXPECT_NEAR(dlogj(i), x_adj(i), tol(dlogj(i))) + << "dlogJ/dx at x = " << x(i); + } + } + stan::math::recover_memory(); + } +} + #endif diff --git a/test/unit/math/rev/constraint/lub_constrain_test.cpp b/test/unit/math/rev/constraint/lub_constrain_test.cpp new file mode 100644 index 00000000000..bed0c1ddef4 --- /dev/null +++ b/test/unit/math/rev/constraint/lub_constrain_test.cpp @@ -0,0 +1,240 @@ +#include +#include +#include +#include +#include +#include +#include + +namespace lub_constrain_rev_test { +// Reference derivatives of y = lb + (ub - lb) * inv_logit(x) and of the +// log Jacobian log(ub - lb) + log(inv_logit(x)) + log(inv_logit(-x)) at +// lb = -2, ub = 3, computed with mpmath at 50 digits: +// dy/dx = (ub - lb) * inv_logit(x) * inv_logit(-x), dy/dlb = inv_logit(-x), +// dy/dub = inv_logit(x), dlogJ/dx = inv_logit(-x) - inv_logit(x), +// dlogJ/dlb = -1 / (ub - lb), dlogJ/dub = 1 / (ub - lb). +// Entries below the smallest subnormal double are 0. +const std::vector xs{-800, -40, -37, -30, -20, -1, 0, + 1, 20, 30, 37, 40, 800}; +const std::vector dx{0.0, + 2.1241771276457945e-17, + 4.2665238128720322e-16, + 4.6788114844192117e-13, + 1.0305768069709247e-8, + 9.8305966620740926e-1, + 1.25, + 9.8305966620740926e-1, + 1.0305768069709247e-8, + 4.6788114844192117e-13, + 4.2665238128720322e-16, + 2.1241771276457945e-17, + 0.0}; +const std::vector dlb{1.0, + 1.0, + 9.9999999999999991e-1, + 9.9999999999990642e-1, + 9.9999999793884638e-1, + 7.3105857863000488e-1, + 0.5, + 2.6894142136999512e-1, + 2.0611536181902036e-9, + 9.357622968839299e-14, + 8.5330476257440651e-17, + 4.248354255291589e-18, + 0.0}; +const std::vector dub{0.0, + 4.248354255291589e-18, + 8.5330476257440651e-17, + 9.357622968839299e-14, + 2.0611536181902036e-9, + 2.6894142136999512e-1, + 0.5, + 7.3105857863000488e-1, + 9.9999999793884638e-1, + 9.9999999999990642e-1, + 9.9999999999999991e-1, + 1.0, + 1.0}; +const std::vector dlogj{1.0, + 9.9999999999999999e-1, + 9.9999999999999983e-1, + 9.9999999999981285e-1, + 9.9999999587769276e-1, + 4.6211715726000976e-1, + 0.0, + -4.6211715726000976e-1, + -9.9999999587769276e-1, + -9.9999999999981285e-1, + -9.9999999999999983e-1, + -9.9999999999999999e-1, + -1.0}; +constexpr double lb = -2.0; +constexpr double ub = 3.0; + +// Relative check, absolute below the smallest normal double. +inline void expect_rel(double expected, double actual, const char* what, + size_t i) { + EXPECT_NEAR( + expected, actual, + 1e-13 + * std::fmax(std::fabs(expected), std::numeric_limits::min())) + << what << " at x = " << xs[i]; +} + +// Checks the adjoint of the first element (or the scalar) of a. +template +inline void expect_adj(const T& a, double expected, const char* what, + size_t i) { + if constexpr (std::is_same::value) { + expect_rel(expected, a.adj(), what, i); + } else if constexpr (stan::is_autodiff_v) { + expect_rel(expected, a.adj()(0), what, i); + } +} + +template +inline void scalar_case(size_t i) { + { + TX x = xs[i]; + TL l = lb; + TU u = ub; + stan::math::var y = stan::math::lub_constrain(x, l, u); + y.grad(); + expect_adj(x, dx[i], "dx", i); + expect_adj(l, dlb[i], "dlb", i); + expect_adj(u, dub[i], "dub", i); + stan::math::recover_memory(); + } + TX x = xs[i]; + TL l = lb; + TU u = ub; + stan::math::var lp = 0; + stan::math::var y = stan::math::lub_constrain(x, l, u, lp); + y.grad(); + expect_adj(x, dx[i], "dx, lp variant", i); + expect_adj(l, dlb[i], "dlb, lp variant", i); + expect_adj(u, dub[i], "dub, lp variant", i); + stan::math::set_zero_all_adjoints(); + lp.grad(); + expect_adj(x, dlogj[i], "dlogJ/dx", i); + expect_adj(l, -0.2, "dlogJ/dlb", i); + expect_adj(u, 0.2, "dlogJ/dub", i); + stan::math::recover_memory(); +} + +// x, lb, ub are vectors or scalars; the first element holds xs[i]. With +// tail, a second element with an infinite vector bound takes the +// overloads' mixed-infinity path. +template +inline T make(const Eigen::VectorXd& v, double s) { + if constexpr (stan::is_stan_scalar::value) { + return T(s); + } else { + return T(v); + } +} + +template +inline void vector_case(size_t i, bool tail) { + using stan::math::var; + constexpr bool l_vec = !stan::is_stan_scalar::value; + constexpr bool u_vec = !stan::is_stan_scalar::value; + const int n = tail ? 2 : 1; + Eigen::VectorXd x_d(n), l_d(n), u_d(n); + x_d(0) = xs[i]; + l_d(0) = lb; + u_d(0) = ub; + if (tail) { + x_d(1) = 0.3; + l_d(1) = l_vec ? stan::math::NEGATIVE_INFTY : lb; + u_d(1) = u_vec && !l_vec ? stan::math::INFTY : ub; + } + for (bool use_lp : {false, true}) { + TX x = x_d; + TL l = make(l_d, lb); + TU u = make(u_d, ub); + var lp = 0; + var s = use_lp ? stan::math::sum(stan::math::lub_constrain(x, l, u, lp)) + : stan::math::sum(stan::math::lub_constrain(x, l, u)); + s.grad(); + expect_adj(x, dx[i], "dx", i); + if (l_vec || !tail) { + expect_adj(l, dlb[i], "dlb", i); + } + if (u_vec || !tail) { + expect_adj(u, dub[i], "dub", i); + } + if (use_lp) { + stan::math::set_zero_all_adjoints(); + lp.grad(); + expect_adj(x, dlogj[i], "dlogJ/dx", i); + if (l_vec || !tail) { + expect_adj(l, -0.2, "dlogJ/dlb", i); + } + if (u_vec || !tail) { + expect_adj(u, 0.2, "dlogJ/dub", i); + } + } + stan::math::recover_memory(); + } +} + +template +inline void vector_cases(size_t i) { + using stan::math::var; + for (bool tail : {false, true}) { + vector_case(i, tail); + vector_case(i, tail); + vector_case(i, tail); + vector_case(i, tail); + vector_case(i, tail); + vector_case(i, tail); + vector_case(i, tail); + } + vector_case(i, false); + vector_case(i, false); + vector_case(i, false); +} +} // namespace lub_constrain_rev_test + +TEST_F(AgradRev, lub_constrain_scalar_derivatives_tails) { + using lub_constrain_rev_test::scalar_case; + using stan::math::var; + for (size_t i = 0; i < lub_constrain_rev_test::xs.size(); ++i) { + scalar_case(i); + scalar_case(i); + scalar_case(i); + scalar_case(i); + scalar_case(i); + } +} + +TEST_F(AgradRev, lub_constrain_matrix_derivatives_tails) { + using stan::math::var; + for (size_t i = 0; i < lub_constrain_rev_test::xs.size(); ++i) { + lub_constrain_rev_test::vector_cases>(i); + } +} + +TEST_F(AgradRev, lub_constrain_var_matrix_derivatives_tails) { + using stan::math::var_value; + for (size_t i = 0; i < lub_constrain_rev_test::xs.size(); ++i) { + lub_constrain_rev_test::vector_cases>(i); + } +} + +TEST_F(AgradRev, lub_constrain_std_vector_derivatives_tails) { + using stan::math::var; + for (size_t i = 0; i < lub_constrain_rev_test::xs.size(); ++i) { + std::vector x{lub_constrain_rev_test::xs[i]}; + var l = lub_constrain_rev_test::lb; + var u = lub_constrain_rev_test::ub; + std::vector y = stan::math::lub_constrain(x, l, u); + y[0].grad(); + lub_constrain_rev_test::expect_adj(x[0], lub_constrain_rev_test::dx[i], + "dx", i); + lub_constrain_rev_test::expect_adj(l, lub_constrain_rev_test::dlb[i], "dlb", + i); + stan::math::recover_memory(); + } +}