From 54cc13312ae7dfe83da6ce13d9038e3e888efc07 Mon Sep 17 00:00:00 2001 From: "Jachym.Barvinek" Date: Mon, 12 Oct 2026 11:25:49 +0200 Subject: [PATCH 1/7] Tests: lub_constrain derivatives at large arguments --- .../rev/constraint/lub_constrain_test.cpp | 54 ++++ .../rev/constraint/lub_constrain_test.cpp | 235 ++++++++++++++++++ 2 files changed, 289 insertions(+) create mode 100644 test/unit/math/rev/constraint/lub_constrain_test.cpp 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..7d88eb1c03b 100644 --- a/test/unit/math/opencl/rev/constraint/lub_constrain_test.cpp +++ b/test/unit/math/opencl/rev/constraint/lub_constrain_test.cpp @@ -125,4 +125,58 @@ 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); + 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), 1e-13 * dx(i)) << "dx at x = " << x(i); + EXPECT_NEAR(dlb(i), lb_adj(i), 1e-13 * 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), 1e-13 * std::fabs(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..6d0daea3dc7 --- /dev/null +++ b/test/unit/math/rev/constraint/lub_constrain_test.cpp @@ -0,0 +1,235 @@ +#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; + +inline void expect_rel(double expected, double actual, const char* what, + size_t i) { + EXPECT_NEAR(expected, actual, 1e-13 * std::fabs(expected)) + << 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(); + } +} From 915a84e03e785ff8db6af0990fe6ad50db4b8df8 Mon Sep 17 00:00:00 2001 From: "Jachym.Barvinek" Date: Mon, 12 Oct 2026 11:25:50 +0200 Subject: [PATCH 2/7] Accurate lub_constrain derivatives for large positive arguments --- .../opencl/rev/constraint/lub_constrain.hpp | 6 +- stan/math/rev/constraint/lub_constrain.hpp | 186 +++++++++++------- 2 files changed, 114 insertions(+), 78 deletions(-) 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/rev/constraint/lub_constrain.hpp b/stan/math/rev/constraint/lub_constrain.hpp index df5d9f9dd46..76b9ce71b00 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,19 @@ namespace stan { namespace math { +namespace internal { +/** + * Return in arena memory the inverse logit of x given exp(x). + * + * @tparam T type of the array of exp(x) + * @param exp_x exp(x) + * @return inverse logit of x + */ +template +inline auto lub_inv_logit_of_exp(const T& exp_x) { + return to_arena((exp_x == INFTY).select(1.0, exp_x / (1.0 + exp_x))); +} +} // namespace internal /** * Return the lower and upper-bounded scalar derived by @@ -48,20 +63,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 +135,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 +185,23 @@ 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_x = to_arena(arena_x.val().array().exp()); + auto inv_logit_x = internal::lub_inv_logit_of_exp(exp_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_x, inv_logit_x]() mutable { + const auto inv_logit_neg_x = (1.0 + exp_x).inverse(); + 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); } } @@ -209,24 +232,26 @@ inline auto lub_constrain(const T& x, const L& lb, const U& ub, 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_x = to_arena(value_of(arena_x).array().exp()); + auto inv_logit_x = internal::lub_inv_logit_of_exp(exp_x); 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_x, inv_logit_x]() mutable { + const auto inv_logit_neg_x = (1.0 + exp_x).inverse(); 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 +284,17 @@ 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_x = to_arena(value_of(arena_x).array().exp()); + auto inv_logit_x = internal::lub_inv_logit_of_exp(exp_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, + reverse_pass_callback([arena_x, ub, arena_lb, ret, diff, exp_x, inv_logit_x, is_lb_inf]() mutable { + const auto inv_logit_neg_x = (1.0 + exp_x).inverse(); 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 +304,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); @@ -309,7 +336,8 @@ inline auto lub_constrain(const T& x, const L& lb, const U& ub, 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_x = to_arena(arena_x_val.exp()); + auto inv_logit_x = internal::lub_inv_logit_of_exp(exp_x); arena_t ret = (is_lb_inf).select(ub_val - arena_x_val.exp(), diff * inv_logit_x + lb_val); lp += (is_lb_inf) @@ -317,19 +345,19 @@ inline auto lub_constrain(const T& x, const L& lb, const U& ub, log(diff) + (neg_abs_x - (2.0 * log1p_exp(neg_abs_x)))) .sum(); reverse_pass_callback([arena_x, arena_x_val, ub, arena_lb, ret, lp, diff, - inv_logit_x, is_lb_inf]() mutable { + exp_x, inv_logit_x, is_lb_inf]() mutable { + const auto inv_logit_neg_x = (1.0 + exp_x).inverse(); 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 +394,22 @@ 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_x = to_arena(arena_x_val.array().exp()); + auto inv_logit_x = internal::lub_inv_logit_of_exp(exp_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_x, inv_logit_x, diff]() mutable { + const auto inv_logit_neg_x = (1.0 + exp_x).inverse(); 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) { @@ -420,22 +450,24 @@ inline auto lub_constrain(const T& x, const L& lb, const U& ub, .select(arena_x_val.array(), log(diff) + (neg_abs_x - (2.0 * log1p_exp(neg_abs_x)))) .sum(); - auto inv_logit_x = to_arena(inv_logit(arena_x_val.array())); + auto exp_x = to_arena(arena_x_val.array().exp()); + auto inv_logit_x = internal::lub_inv_logit_of_exp(exp_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, diff, inv_logit_x, arena_ub, - lb, ret, lp, is_ub_inf]() mutable { + reverse_pass_callback([arena_x, arena_x_val, diff, exp_x, inv_logit_x, + arena_ub, lb, ret, lp, is_ub_inf]() mutable { + const auto inv_logit_neg_x = (1.0 + exp_x).inverse(); 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 +496,8 @@ 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_x = to_arena(arena_x_val.array().exp()); + auto inv_logit_x = internal::lub_inv_logit_of_exp(exp_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 +510,19 @@ 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_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 = (1.0 + exp_x).inverse(); // 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 +538,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 +548,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 +572,8 @@ 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_x = to_arena(arena_x_val.array().exp()); + auto inv_logit_x = internal::lub_inv_logit_of_exp(exp_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); @@ -560,21 +595,22 @@ inline auto lub_constrain(const T& x, const L& lb, const U& ub, arena_x_val.array(), log(diff) + (neg_abs_x - (2.0 * log1p_exp(neg_abs_x))))) .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_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 = (1.0 + exp_x).inverse(); 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 +629,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) { From 8cae6d4630058bb702ca6ca8ca6c8a304ac480c0 Mon Sep 17 00:00:00 2001 From: Aki Vehtari Date: Sun, 18 Oct 2026 23:44:36 +0300 Subject: [PATCH 3/7] Compute inv_logit(-x) without a division in the lub_constrain matrix reverse passes --- stan/math/rev/constraint/lub_constrain.hpp | 131 +++++++++++++-------- 1 file changed, 80 insertions(+), 51 deletions(-) diff --git a/stan/math/rev/constraint/lub_constrain.hpp b/stan/math/rev/constraint/lub_constrain.hpp index 76b9ce71b00..a265e77585d 100644 --- a/stan/math/rev/constraint/lub_constrain.hpp +++ b/stan/math/rev/constraint/lub_constrain.hpp @@ -14,15 +14,21 @@ namespace stan { namespace math { namespace internal { /** - * Return in arena memory the inverse logit of x given exp(x). + * Return in arena memory the inverse logit of x given exp(-|x|). * - * @tparam T type of the array of exp(x) - * @param exp_x exp(x) + * The reverse passes compute inv_logit(-x) from the result without a + * division: exp(-|x|) * inv_logit(x) for x >= 0, and 1 - inv_logit(x), + * which is exact for inv_logit(x) <= 1/2, otherwise. + * + * @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_of_exp(const T& exp_x) { - return to_arena((exp_x == INFTY).select(1.0, exp_x / (1.0 + exp_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 @@ -185,12 +191,15 @@ 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 exp_x = to_arena(arena_x.val().array().exp()); - auto inv_logit_x = internal::lub_inv_logit_of_exp(exp_x); + 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, exp_x, inv_logit_x]() mutable { - const auto inv_logit_neg_x = (1.0 + exp_x).inverse(); + [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; @@ -231,13 +240,16 @@ inline auto lub_constrain(const T& x, const L& lb, const U& ub, 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 exp_x = to_arena(value_of(arena_x).array().exp()); - auto inv_logit_x = internal::lub_inv_logit_of_exp(exp_x); + 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); + lp += (log(diff) + (neg_abs_x - (2.0 * exp_neg_abs_x.log1p()))).sum(); arena_t ret = diff * inv_logit_x + lb_val; reverse_pass_callback( - [arena_x, ub, lb, ret, lp, diff, exp_x, inv_logit_x]() mutable { - const auto inv_logit_neg_x = (1.0 + exp_x).inverse(); + [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 * inv_logit_neg_x @@ -284,13 +296,16 @@ 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 exp_x = to_arena(value_of(arena_x).array().exp()); - auto inv_logit_x = internal::lub_inv_logit_of_exp(exp_x); + 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, exp_x, inv_logit_x, - is_lb_inf]() mutable { - const auto inv_logit_neg_x = (1.0 + exp_x).inverse(); + 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(), @@ -336,17 +351,19 @@ inline auto lub_constrain(const T& x, const L& lb, const U& ub, 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 exp_x = to_arena(arena_x_val.exp()); - auto inv_logit_x = internal::lub_inv_logit_of_exp(exp_x); + 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)))) + log(diff) + (neg_abs_x - (2.0 * exp_neg_abs_x.log1p()))) .sum(); reverse_pass_callback([arena_x, arena_x_val, ub, arena_lb, ret, lp, diff, - exp_x, inv_logit_x, is_lb_inf]() mutable { - const auto inv_logit_neg_x = (1.0 + exp_x).inverse(); + 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(); @@ -394,13 +411,16 @@ 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 exp_x = to_arena(arena_x_val.array().exp()); - auto inv_logit_x = internal::lub_inv_logit_of_exp(exp_x); + 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, - exp_x, inv_logit_x, diff]() mutable { - const auto inv_logit_neg_x = (1.0 + exp_x).inverse(); + 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(), @@ -446,17 +466,21 @@ inline auto lub_constrain(const T& x, const L& lb, const U& ub, 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) + (neg_abs_x - (2.0 * exp_neg_abs_x.log1p()))) .sum(); - auto exp_x = to_arena(arena_x_val.array().exp()); - auto inv_logit_x = internal::lub_inv_logit_of_exp(exp_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, diff, exp_x, inv_logit_x, - arena_ub, lb, ret, lp, is_ub_inf]() mutable { - const auto inv_logit_neg_x = (1.0 + exp_x).inverse(); + 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( @@ -496,8 +520,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 exp_x = to_arena(arena_x_val.array().exp()); - auto inv_logit_x = internal::lub_inv_logit_of_exp(exp_x); + 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); @@ -510,10 +535,12 @@ 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, exp_x, 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 = (1.0 + exp_x).inverse(); + 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) { @@ -572,8 +599,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 exp_x = to_arena(arena_x_val.array().exp()); - auto inv_logit_x = internal::lub_inv_logit_of_exp(exp_x); + 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); @@ -588,17 +616,18 @@ inline auto lub_constrain(const T& x, const L& lb, const U& ub, 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) + + (neg_abs_x + - (2.0 * exp_neg_abs_x.log1p())))) .sum(); - reverse_pass_callback([arena_x, arena_x_val, exp_x, 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, lp]() mutable { - const auto inv_logit_neg_x = (1.0 + exp_x).inverse(); + 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()); From 456e2579eda22755482c7d557524e544264ea6aa Mon Sep 17 00:00:00 2001 From: Aki Vehtari Date: Sun, 18 Oct 2026 23:44:42 +0300 Subject: [PATCH 4/7] Fix forward-mode lub_constrain derivatives for large x and at x = 0 --- stan/math/prim/constraint/lub_constrain.hpp | 31 ++++- .../mix/constraint/lub_constrain_fwd_test.cpp | 117 ++++++++++++++++++ 2 files changed, 145 insertions(+), 3 deletions(-) create mode 100644 test/unit/math/mix/constraint/lub_constrain_fwd_test.cpp 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/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(); + } +} From 52828dd7027f7d6cd48ce26d54c5bf8dfaed7a22 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?J=C3=A1chym=20Barv=C3=ADnek?= Date: Thu, 22 Oct 2026 01:32:18 +0200 Subject: [PATCH 5/7] Tests: floor the lub_constrain tolerances at the smallest normal double The rev and OpenCL tail tests used 1e-13 * |expected| as the tolerance, so a zero reference (x = +-800, dlogJ/dx at x = 0) demanded an exact 0. Use 1e-13 * max(|expected|, DBL_MIN), as the forward-mode test does. --- .../opencl/rev/constraint/lub_constrain_test.cpp | 13 ++++++++++--- .../unit/math/rev/constraint/lub_constrain_test.cpp | 7 ++++++- 2 files changed, 16 insertions(+), 4 deletions(-) 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 7d88eb1c03b..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); @@ -149,6 +151,11 @@ TEST(OpenCLLubConstrain, rev_derivatives_tails) { -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); @@ -163,15 +170,15 @@ TEST(OpenCLLubConstrain, rev_derivatives_tails) { 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), 1e-13 * dx(i)) << "dx at x = " << x(i); - EXPECT_NEAR(dlb(i), lb_adj(i), 1e-13 * dlb(i)) << "dlb at x = " << x(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), 1e-13 * std::fabs(dlogj(i))) + EXPECT_NEAR(dlogj(i), x_adj(i), tol(dlogj(i))) << "dlogJ/dx at x = " << x(i); } } diff --git a/test/unit/math/rev/constraint/lub_constrain_test.cpp b/test/unit/math/rev/constraint/lub_constrain_test.cpp index 6d0daea3dc7..bed0c1ddef4 100644 --- a/test/unit/math/rev/constraint/lub_constrain_test.cpp +++ b/test/unit/math/rev/constraint/lub_constrain_test.cpp @@ -2,6 +2,7 @@ #include #include #include +#include #include #include @@ -70,9 +71,13 @@ const std::vector dlogj{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::fabs(expected)) + EXPECT_NEAR( + expected, actual, + 1e-13 + * std::fmax(std::fabs(expected), std::numeric_limits::min())) << what << " at x = " << xs[i]; } From 880c7ce355ae6f4a27b1c6f047156e55c877531d Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?J=C3=A1chym=20Barv=C3=ADnek?= Date: Thu, 22 Oct 2026 01:32:18 +0200 Subject: [PATCH 6/7] Keep -|x| out of the arena in the lub_constrain lp overloads The four lp matrix overloads stored -|x| in the arena although only the forward-pass log Jacobian reads it, so their arena use was 8 bytes per element above develop. Compute it in that expression instead. In the scalar-bound overload sum log1p(e) and |x| separately: with the AoS view of x in the same expression Eigen does not vectorize log1p, which made that path 1.3 to 1.55 times slower with MinGW. Also drop an unused sign() evaluation in a reverse pass and state the helper's switching condition precisely in its comment. --- stan/math/rev/constraint/lub_constrain.hpp | 24 +++++++++++----------- 1 file changed, 12 insertions(+), 12 deletions(-) diff --git a/stan/math/rev/constraint/lub_constrain.hpp b/stan/math/rev/constraint/lub_constrain.hpp index a265e77585d..84850876351 100644 --- a/stan/math/rev/constraint/lub_constrain.hpp +++ b/stan/math/rev/constraint/lub_constrain.hpp @@ -17,8 +17,8 @@ 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) for x >= 0, and 1 - inv_logit(x), - * which is exact for inv_logit(x) <= 1/2, otherwise. + * 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|) @@ -238,12 +238,13 @@ 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; 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); - lp += (log(diff) + (neg_abs_x - (2.0 * exp_neg_abs_x.log1p()))).sum(); + // 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, exp_neg_abs_x, inv_logit_x]() mutable { @@ -350,14 +351,14 @@ 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 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 * exp_neg_abs_x.log1p()))) + .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, exp_neg_abs_x, inv_logit_x, is_lb_inf]() mutable { @@ -366,7 +367,6 @@ inline auto lub_constrain(const T& x, const L& lb, const U& ub, .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 * inv_logit_neg_x @@ -465,13 +465,14 @@ 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 * exp_neg_abs_x.log1p()))) + log(diff) + + (-(arena_x_val.array()).abs() + - (2.0 * exp_neg_abs_x.log1p()))) .sum(); arena_t ret = (is_ub_inf).select( arena_x_val.array().exp() + lb_val, diff * inv_logit_x + lb_val); @@ -614,12 +615,11 @@ 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 + + (-(arena_x_val.array()).abs() - (2.0 * exp_neg_abs_x.log1p())))) .sum(); reverse_pass_callback([arena_x, arena_x_val, exp_neg_abs_x, inv_logit_x, From 527ab51bba8ffa3a4f78da12933d3c73028b8591 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?J=C3=A1chym=20Barv=C3=ADnek?= Date: Thu, 22 Oct 2026 10:58:36 +0200 Subject: [PATCH 7/7] Move the value_of_rec include of prim lub_constrain next to Eigen.hpp At the end of the list it sat next to sum.hpp, which #3470 removes as unused, so the two PRs conflicted on that line. Placed after Eigen.hpp it merges cleanly with #3470. --- stan/math/prim/constraint/lub_constrain.hpp | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/stan/math/prim/constraint/lub_constrain.hpp b/stan/math/prim/constraint/lub_constrain.hpp index 340f7020b87..d785a3a9367 100644 --- a/stan/math/prim/constraint/lub_constrain.hpp +++ b/stan/math/prim/constraint/lub_constrain.hpp @@ -5,6 +5,7 @@ #include #include #include +#include #include #include #include @@ -17,7 +18,6 @@ #include #include #include -#include #include namespace stan {