diff --git a/stan/math/prim/fun/dims.hpp b/stan/math/prim/fun/dims.hpp index 332a32e6089..f85a7e78410 100644 --- a/stan/math/prim/fun/dims.hpp +++ b/stan/math/prim/fun/dims.hpp @@ -45,7 +45,11 @@ template inline void dims(const std::vector& x, std::vector& result) { result.push_back(x.size()); if (x.size() > 0) { - dims(x[0], result); + if constexpr (is_eigen_vector::value) { + result.push_back(x[0].size()); + } else { + dims(x[0], result); + } } } diff --git a/stan/math/prim/functor.hpp b/stan/math/prim/functor.hpp index 22a3934f99d..782a87c5af8 100644 --- a/stan/math/prim/functor.hpp +++ b/stan/math/prim/functor.hpp @@ -37,6 +37,7 @@ #include #include #include +#include #include #endif diff --git a/stan/math/prim/functor/solve_itp.hpp b/stan/math/prim/functor/solve_itp.hpp new file mode 100644 index 00000000000..a2d6869a196 --- /dev/null +++ b/stan/math/prim/functor/solve_itp.hpp @@ -0,0 +1,311 @@ +#ifndef STAN_MATH_PRIM_FUNCTOR_SOLVE_ITP_HPP +#define STAN_MATH_PRIM_FUNCTOR_SOLVE_ITP_HPP + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +namespace stan { +namespace math { +namespace internal { + +/** + * Return ceil(log2(width / (2 * x_tolerance))), truncated below at zero, + * without forming a potentially overflowing or underflowing ratio. + */ +inline int64_t itp_bisection_iterations(double width, double x_tolerance) { + if (width <= 2.0 * x_tolerance) { + return 0; + } + + int width_exponent; + int tolerance_exponent; + const double width_mantissa = std::frexp(width, &width_exponent); + const double tolerance_mantissa + = std::frexp(x_tolerance, &tolerance_exponent); + + // ceil(log2(width / x_tolerance)) is the exponent difference plus one + // exactly when the mantissa ratio is greater than one. Subtracting one + // then gives ceil(log2(width / (2 * x_tolerance))). + const int64_t ceil_log2_ratio + = static_cast(width_exponent) + - static_cast(tolerance_exponent) + + (width_mantissa > tolerance_mantissa ? 1 : 0); + return std::max(0, ceil_log2_ratio - 1); +} + +/** + * Scale x_tolerance by an integral power of two. The result is allowed to + * saturate at positive infinity because it is subsequently used only as an + * upper projection radius. + */ +inline double itp_scaled_tolerance(double x_tolerance, int64_t exponent) { + if (exponent > std::numeric_limits::max()) { + return std::numeric_limits::infinity(); + } + if (exponent < std::numeric_limits::min()) { + return 0.0; + } + return std::ldexp(x_tolerance, static_cast(exponent)); +} + +/** + * Return the default truncation scale without overflowing on an interval whose + * positive width is subnormal. The fallback value is irrelevant when the + * default tolerance already contains the entire interval, but keeping it finite + * allows endpoint and bracket checks to retain their normal semantics. + */ +inline double itp_default_kappa1(double width) { + const double overflow_threshold = 0.2 / std::numeric_limits::max(); + return width > overflow_threshold ? 0.2 / width : 1.0; +} + +/** + * Return the regula-falsi point for an opposite-signed bracket. This form is + * algebraically equivalent to + * + * (f_b * a - f_a * b) / (f_b - f_a), + * + * but avoids products of a coordinate and a potentially extreme function + * value, and avoids overflow in |f_a| + |f_b|. + */ +inline double itp_false_position(double a, double b, double f_a, double f_b) { + const double abs_f_a = std::abs(f_a); + const double abs_f_b = std::abs(f_b); + double lambda; + + if (abs_f_a <= abs_f_b) { + const double ratio = abs_f_a / abs_f_b; + lambda = ratio / (1.0 + ratio); + } else { + const double ratio = abs_f_b / abs_f_a; + lambda = 1.0 / (1.0 + ratio); + } + return a + lambda * (b - a); +} + +inline bool itp_opposite_signs(double x, double y) { + return (x < 0.0 && y > 0.0) || (x > 0.0 && y < 0.0); +} + +inline void itp_check_controls(const char* function, double lower, double upper, + double x_tolerance, double kappa1, double kappa2, + int64_t n0, int64_t max_num_steps) { + static constexpr double one_plus_phi + = 0.5 * (3.0 + 2.23606797749978969640917366873127623544); + + check_finite(function, "lower bound", lower); + check_finite(function, "upper bound", upper); + check_less(function, "lower bound", lower, upper); + check_finite(function, "x tolerance", x_tolerance); + check_positive(function, "x tolerance", x_tolerance); + check_finite(function, "kappa1", kappa1); + check_positive(function, "kappa1", kappa1); + check_finite(function, "kappa2", kappa2); + check_greater_or_equal(function, "kappa2", kappa2, 1.0); + check_less(function, "kappa2", kappa2, one_plus_phi); + check_nonnegative(function, "n0", n0); + check_positive(function, "max_num_steps", max_num_steps); + + const double width = upper - lower; + check_finite(function, "upper bound minus lower bound", width); +} + +/** + * Primitive ITP implementation after validation of control arguments. + */ +template +inline double solve_itp_primitive(const F& f, double lower, double upper, + double x_tolerance, double kappa1, + double kappa2, int64_t n0, + int64_t max_num_steps, std::ostream* msgs, + const Args&... args) { + static constexpr const char* function = "solve_itp"; + double a = lower; + double b = upper; + double f_a = value_of(f(a, msgs, args...)); + double f_b = value_of(f(b, msgs, args...)); + + check_finite(function, "function value at lower bound", f_a); + check_finite(function, "function value at upper bound", f_b); + + if (f_a == 0.0) { + return a; + } + if (f_b == 0.0) { + return b; + } + if (!itp_opposite_signs(f_a, f_b)) { + std::ostringstream message; + message << function + << ": function values at the lower and upper bounds are " << f_a + << " and " << f_b << ", but must have opposite signs"; + throw std::domain_error(message.str()); + } + + const int64_t n_half = itp_bisection_iterations(b - a, x_tolerance); + if (n0 > std::numeric_limits::max() - n_half) { + throw_domain_error(function, "n0", n0, "", + " is too large to form the ITP iteration bound"); + } + const int64_t n_max = n_half + n0; + if (n_max > max_num_steps) { + throw_domain_error(function, "theoretical ITP iteration bound", n_max, "", + " exceeds max_num_steps"); + } + + int64_t iteration = 0; + const double two_tolerance = 2.0 * x_tolerance; + while (b - a > two_tolerance) { + // The real-arithmetic proof gives n_max iterations. One additional + // midpoint step is permitted solely to repair an outward-rounded final + // width such as nextafter(2 * x_tolerance, +infinity). + if (iteration > n_max) { + throw_domain_error( + function, "number of ITP iterations", iteration, "", + " exceeded the exact-arithmetic bound plus one roundoff repair"); + } + if (iteration >= max_num_steps) { + throw_domain_error(function, "number of ITP iterations", iteration, "", + " reached max_num_steps before convergence"); + } + + const double width = b - a; + const double midpoint = 0.5 * a + 0.5 * b; + if (!(midpoint > a && midpoint < b)) { + throw_domain_error( + function, "x tolerance", x_tolerance, "", + " is smaller than the representable spacing in the final bracket"); + } + + double x_itp = midpoint; + if (iteration < n_max) { + // Interpolation. + const double x_f = itp_false_position(a, b, f_a, f_b); + + // Truncation toward the midpoint. + const double gap = std::abs(midpoint - x_f); + const double sigma = midpoint > x_f ? 1.0 : (midpoint < x_f ? -1.0 : 0.0); + double x_t = midpoint; + if (gap > 0.0) { + const double log_delta = std::log(kappa1) + kappa2 * std::log(width); + if (log_delta < std::log(gap)) { + x_t = x_f + sigma * std::exp(log_delta); + } + } + + // Projection onto the minmax-safe interval centered at the midpoint. + const double scaled_tolerance + = itp_scaled_tolerance(x_tolerance, n_max - iteration); + double projection_radius = scaled_tolerance - 0.5 * width; + projection_radius = std::max(0.0, projection_radius); + const double projected_distance + = std::min(projection_radius, std::abs(x_t - midpoint)); + x_itp = midpoint - sigma * projected_distance; + + // Midpoint fallback protects the bracket from exceptional roundoff. + if (!std::isfinite(x_itp) || !(x_itp > a && x_itp < b)) { + x_itp = midpoint; + } + } + + const double f_itp = value_of(f(x_itp, msgs, args...)); + check_finite(function, "function value at candidate point", f_itp); + if (f_itp == 0.0) { + return x_itp; + } + + if (itp_opposite_signs(f_a, f_itp)) { + b = x_itp; + f_b = f_itp; + } else { + a = x_itp; + f_a = f_itp; + } + ++iteration; + } + + return 0.5 * a + 0.5 * b; +} + +} // namespace internal + +/** + * Find a root of a continuous scalar function in a finite sign-changing + * bracket using the interpolate-truncate-project (ITP) method. + * + * The callback must be callable as + * + * f(x, msgs, args...) + * + * and return a scalar. The interval endpoints and tuning controls are data; + * differentiation with respect to args is supplied by the reverse-mode + * overload. + * + * @tparam F callback type + * @tparam Args types of additional callback arguments + * @param f scalar callback + * @param lower lower endpoint of the initial bracket + * @param upper upper endpoint of the initial bracket + * @param x_tolerance requested absolute error bound on the returned midpoint + * @param kappa1 positive truncation scale + * @param kappa2 truncation exponent in [1, 1 + phi) + * @param n0 nonnegative slack in the bisection minmax iteration bound + * @param max_num_steps maximum permitted number of function-evaluation + * iterations; one step beyond the exact-arithmetic ITP bound may be + * needed to repair an outward-rounded final width + * @param msgs stream for callback messages + * @param args additional callback arguments + * @return midpoint of a final sign-changing bracket whose computed width is + * at most 2 * x_tolerance, or an exactly evaluated root + */ +template * = nullptr> +inline double solve_itp_tol(const F& f, double lower, double upper, + double x_tolerance, double kappa1, double kappa2, + int64_t n0, int64_t max_num_steps, + std::ostream* msgs, const Args&... args) { + static constexpr const char* function = "solve_itp"; + internal::itp_check_controls(function, lower, upper, x_tolerance, kappa1, + kappa2, n0, max_num_steps); + return internal::solve_itp_primitive(f, lower, upper, x_tolerance, kappa1, + kappa2, n0, max_num_steps, msgs, + args...); +} + +/** + * Find a root using the default ITP controls + * + * x_tolerance = 1e-10, + * kappa1 = 0.2 / (upper - lower), + * kappa2 = 2, + * n0 = 1, + * max steps = 200. + */ +template * = nullptr> +inline double solve_itp(const F& f, double lower, double upper, + std::ostream* msgs, const Args&... args) { + static constexpr const char* function = "solve_itp"; + check_finite(function, "lower bound", lower); + check_finite(function, "upper bound", upper); + check_less(function, "lower bound", lower, upper); + const double width = upper - lower; + check_finite(function, "upper bound minus lower bound", width); + + return solve_itp_tol(f, lower, upper, 1e-10, + internal::itp_default_kappa1(width), 2.0, 1, 200, msgs, + args...); +} + +} // namespace math +} // namespace stan + +#endif diff --git a/stan/math/rev/functor.hpp b/stan/math/rev/functor.hpp index 6539fcaf5ba..0615b9bb779 100644 --- a/stan/math/rev/functor.hpp +++ b/stan/math/rev/functor.hpp @@ -4,6 +4,7 @@ #include #include #include +#include #include #include #include diff --git a/stan/math/rev/functor/solve_itp.hpp b/stan/math/rev/functor/solve_itp.hpp new file mode 100644 index 00000000000..41dbeb2c97a --- /dev/null +++ b/stan/math/rev/functor/solve_itp.hpp @@ -0,0 +1,108 @@ +#ifndef STAN_MATH_REV_FUNCTOR_SOLVE_ITP_HPP +#define STAN_MATH_REV_FUNCTOR_SOLVE_ITP_HPP + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +namespace stan { +namespace math { + +/** + * Reverse-mode overload of the ITP scalar root solver. + * + * The root is found with primitive values. Gradients with respect to the + * callback arguments are then supplied by the scalar implicit-function rule + * + * dx / dtheta = -(partial f / partial theta) / (partial f / partial x) + * + * evaluated at the numerical root. The bracket and all solver controls are + * treated as data and are not differentiated. + */ +template * = nullptr> +inline var solve_itp_tol(const F& f, double lower, double upper, + double x_tolerance, double kappa1, double kappa2, + int64_t n0, int64_t max_num_steps, std::ostream* msgs, + const Args&... args) { + static constexpr const char* function = "solve_itp"; + + auto arena_args_tuple = make_chainable_ptr(std::make_tuple(eval(args)...)); + auto args_vals_tuple = math::apply( + [&](const auto&... arena_args) { + return std::make_tuple(to_ref(value_of(arena_args))...); + }, + *arena_args_tuple); + + const double root = math::apply( + [&](const auto&... args_vals) { + return solve_itp_tol(f, lower, upper, x_tolerance, kappa1, kappa2, n0, + max_num_steps, msgs, args_vals...); + }, + args_vals_tuple); + + // Compute partial f / partial x with parameters held at their values. + double f_x; + { + nested_rev_autodiff nested; + var x_var = root; + var residual = math::apply( + [&](const auto&... args_vals) { + return eval(f(x_var, msgs, args_vals...)); + }, + args_vals_tuple); + residual.grad(); + f_x = x_var.adj(); + } + + check_finite(function, "partial derivative with respect to the root", f_x); + if (f_x == 0.0) { + throw_domain_error(function, "partial derivative with respect to the root", + f_x, "", + " is zero, so the implicit derivative is undefined"); + } + + var result = root; + reverse_pass_callback([f, result, arena_args_tuple, f_x, msgs]() mutable { + nested_rev_autodiff nested; + var residual = math::apply( + [&](const auto&... arena_args) { + return eval(f(result.val(), msgs, arena_args...)); + }, + *arena_args_tuple); + residual.adj() = -result.adj() / f_x; + grad(); + }); + + return result; +} + +/** + * Reverse-mode ITP solver with the default controls. + */ +template * = nullptr> +inline var solve_itp(const F& f, double lower, double upper, std::ostream* msgs, + const Args&... args) { + static constexpr const char* function = "solve_itp"; + check_finite(function, "lower bound", lower); + check_finite(function, "upper bound", upper); + check_less(function, "lower bound", lower, upper); + const double width = upper - lower; + check_finite(function, "upper bound minus lower bound", width); + + return solve_itp_tol(f, lower, upper, 1e-10, + internal::itp_default_kappa1(width), 2.0, 1, 200, msgs, + args...); +} + +} // namespace math +} // namespace stan + +#endif diff --git a/test/unit/math/prim/fun/dims_test.cpp b/test/unit/math/prim/fun/dims_test.cpp index 67a988a8598..00e00714a6b 100644 --- a/test/unit/math/prim/fun/dims_test.cpp +++ b/test/unit/math/prim/fun/dims_test.cpp @@ -58,4 +58,32 @@ TEST(MathMatrixPrimMat, dims) { EXPECT_EQ(2, dims8[0]); EXPECT_EQ(7, dims8[1]); EXPECT_EQ(8, dims8[2]); +} + TEST(MathMatrixPrimMat, dims_nested_vector_arrays) { + using Eigen::Dynamic; + using Eigen::Matrix; + using stan::math::dims; + using std::vector; + + Matrix v1(1); + vector> inner1{v1}; + vector>> x1{inner1}; + + const auto result1 = dims(x1); + + ASSERT_EQ(3U, result1.size()); + EXPECT_EQ(1, result1[0]); + EXPECT_EQ(1, result1[1]); + EXPECT_EQ(1, result1[2]); + + Matrix v2(5); + vector> inner2(3, v2); + vector>> x2(2, inner2); + + const auto result2 = dims(x2); + + ASSERT_EQ(3U, result2.size()); + EXPECT_EQ(2, result2[0]); + EXPECT_EQ(3, result2[1]); + EXPECT_EQ(5, result2[2]); } diff --git a/test/unit/math/prim/functor/solve_itp_test.cpp b/test/unit/math/prim/functor/solve_itp_test.cpp new file mode 100644 index 00000000000..fd9445948d6 --- /dev/null +++ b/test/unit/math/prim/functor/solve_itp_test.cpp @@ -0,0 +1,191 @@ +#include +#include +#include +#include +#include + +namespace { + +struct cubic_functor { + mutable int evaluations = 0; + + template + auto operator()(const T& x, std::ostream*, const Args&...) const { + ++evaluations; + return x * x * x - x - 2.0; + } +}; + +struct lambert_functor { + mutable int evaluations = 0; + + template + auto operator()(const T& x, std::ostream*, const Args&...) const { + ++evaluations; + return x * exp(x) - 1.0; + } +}; + +struct affine_functor { + template + auto operator()(const T& x, std::ostream*, const T_scale& scale, + const T_offset& offset) const { + return scale * x - offset; + } +}; + +struct endpoint_functor { + template + auto operator()(const T& x, std::ostream*, const Args&...) const { + return x; + } +}; + +struct no_bracket_functor { + template + auto operator()(const T& x, std::ostream*, const Args&...) const { + return x * x + 1.0; + } +}; + +struct nonfinite_functor { + template + auto operator()(const T& x, std::ostream*, const Args&...) const { + if (x > 0.25 && x < 0.75) { + return T(std::numeric_limits::quiet_NaN()); + } + return x - 0.5; + } +}; + +struct steep_functor { + mutable int evaluations = 0; + + template + auto operator()(const T& x, std::ostream*, const Args&...) const { + ++evaluations; + using std::atan; + return atan(1e12 * (x - 0.123456789)); + } +}; + +struct final_width_roundoff_functor { + mutable int evaluations = 0; + + template + auto operator()(const T& x, std::ostream*, const Args&...) const { + ++evaluations; + constexpr double root = 0.42425901218220274; + constexpr double linear = 0.0004533163653989277; + const auto z = x - root; + return z * z * z + linear * z; + } +}; + +} // namespace + +TEST(MathPrimFunctorSolveItp, BisectionIterationCount) { + using stan::math::internal::itp_bisection_iterations; + EXPECT_EQ(0, itp_bisection_iterations(1.0, 0.5)); + EXPECT_EQ(1, itp_bisection_iterations(1.0, 0.25)); + EXPECT_EQ(3, itp_bisection_iterations(8.0, 0.5)); + EXPECT_EQ(10, itp_bisection_iterations(1.0, 1.0 / 2048.0)); +} + +TEST(MathPrimFunctorSolveItp, PublishedCubicExample) { + cubic_functor f; + const double root + = stan::math::solve_itp_tol(f, 1.0, 2.0, 5e-4, 0.1, 2.0, 1, 100, nullptr); + + EXPECT_NEAR(1.5213797068045676, root, 5e-4); + // Two endpoint evaluations plus five ITP evaluations. + EXPECT_EQ(7, f.evaluations); +} + +TEST(MathPrimFunctorSolveItp, LambertExample) { + lambert_functor f; + const double root = stan::math::solve_itp_tol(f, -1.0, 1.0, 1e-10, 0.1, 2.0, + 1, 100, nullptr); + + EXPECT_NEAR(0.5671432904097838, root, 1e-10); + EXPECT_LE(f.evaluations - 2, 8); +} + +TEST(MathPrimFunctorSolveItp, VariadicArguments) { + const double root + = stan::math::solve_itp(affine_functor(), -10.0, 10.0, nullptr, 3.0, 2.0); + EXPECT_NEAR(2.0 / 3.0, root, 1e-10); +} + +TEST(MathPrimFunctorSolveItp, EndpointRoot) { + EXPECT_EQ(0.0, stan::math::solve_itp(endpoint_functor(), 0.0, 2.0, nullptr)); +} + +TEST(MathPrimFunctorSolveItp, DefaultControlsOnSubnormalWidth) { + const double denorm = std::numeric_limits::denorm_min(); + EXPECT_EQ( + 0.0, stan::math::solve_itp(endpoint_functor(), -denorm, denorm, nullptr)); +} + +TEST(MathPrimFunctorSolveItp, MinmaxEvaluationBound) { + steep_functor f; + constexpr double lower = -1.0; + constexpr double upper = 1.0; + constexpr double x_tolerance = 1e-12; + constexpr int64_t n0 = 0; + const int64_t n_half = stan::math::internal::itp_bisection_iterations( + upper - lower, x_tolerance); + + const double root = stan::math::solve_itp_tol( + f, lower, upper, x_tolerance, 0.1, 2.0, n0, n_half + 1, nullptr); + + EXPECT_NEAR(0.123456789, root, x_tolerance); + EXPECT_LE(f.evaluations, n_half + 3); +} + +TEST(MathPrimFunctorSolveItp, RepairsOutwardRoundedFinalWidth) { + final_width_roundoff_functor f; + constexpr double x_tolerance = 1e-10; + const int64_t exact_bound + = stan::math::internal::itp_bisection_iterations(2.0, x_tolerance) + 1; + + const double root = stan::math::solve_itp_tol( + f, -1.0, 1.0, x_tolerance, 0.1, 2.0, 1, exact_bound + 1, nullptr); + + EXPECT_NEAR(0.42425901218220274, root, x_tolerance); + EXPECT_LE(f.evaluations, exact_bound + 3); +} + +TEST(MathPrimFunctorSolveItp, RejectsNonBracket) { + EXPECT_THROW(stan::math::solve_itp(no_bracket_functor(), -1.0, 1.0, nullptr), + std::domain_error); +} + +TEST(MathPrimFunctorSolveItp, RejectsInvalidControls) { + cubic_functor f; + EXPECT_THROW( + stan::math::solve_itp_tol(f, 1.0, 2.0, 0.0, 0.1, 2.0, 1, 100, nullptr), + std::domain_error); + EXPECT_THROW( + stan::math::solve_itp_tol(f, 1.0, 2.0, 1e-6, 0.0, 2.0, 1, 100, nullptr), + std::domain_error); + EXPECT_THROW( + stan::math::solve_itp_tol(f, 1.0, 2.0, 1e-6, 0.1, 3.0, 1, 100, nullptr), + std::domain_error); + EXPECT_THROW( + stan::math::solve_itp_tol(f, 1.0, 2.0, 1e-6, 0.1, 2.0, -1, 100, nullptr), + std::domain_error); +} + +TEST(MathPrimFunctorSolveItp, RejectsTooFewSteps) { + cubic_functor f; + EXPECT_THROW( + stan::math::solve_itp_tol(f, 1.0, 2.0, 1e-12, 0.1, 2.0, 1, 2, nullptr), + std::domain_error); +} + +TEST(MathPrimFunctorSolveItp, RejectsNonfiniteFunctionValue) { + EXPECT_THROW(stan::math::solve_itp_tol(nonfinite_functor(), 0.0, 1.0, 1e-8, + 0.2, 2.0, 1, 100, nullptr), + std::domain_error); +} diff --git a/test/unit/math/rev/functor/solve_itp_test.cpp b/test/unit/math/rev/functor/solve_itp_test.cpp new file mode 100644 index 00000000000..546be8e2964 --- /dev/null +++ b/test/unit/math/rev/functor/solve_itp_test.cpp @@ -0,0 +1,101 @@ +#include +#include +#include +#include +#include +#include + +namespace { + +struct cubic_parameter_functor { + template + auto operator()(const T_x& x, std::ostream*, const T_theta& theta) const { + return x * x * x - theta; + } +}; + +struct affine_parameter_functor { + template + auto operator()(const T_x& x, std::ostream*, const T_scale& scale, + const T_offset& offset) const { + return scale * x - offset; + } +}; + +struct vector_parameter_functor { + template + auto operator()(const T_x& x, std::ostream*, const T_theta& theta) const { + return theta(0) * x + theta(1) - theta(2); + } +}; + +struct singular_parameter_functor { + template + auto operator()(const T_x& x, std::ostream*, const T_theta& theta) const { + return x * x - theta; + } +}; + +} // namespace + +TEST(MathRevFunctorSolveItp, ImplicitDerivativeScalarParameter) { + using stan::math::var; + var theta = 2.0; + var root = stan::math::solve_itp(cubic_parameter_functor(), 0.0, 2.0, nullptr, + theta); + + std::vector inputs{theta}; + std::vector gradient; + root.grad(inputs, gradient); + + const double expected_root = std::cbrt(2.0); + const double expected_derivative + = 1.0 / (3.0 * expected_root * expected_root); + EXPECT_NEAR(expected_root, root.val(), 1e-10); + EXPECT_NEAR(expected_derivative, gradient[0], 1e-9); + stan::math::recover_memory(); +} + +TEST(MathRevFunctorSolveItp, ImplicitDerivativeVariadicParameters) { + using stan::math::var; + var scale = 3.0; + var offset = 2.0; + var root = stan::math::solve_itp(affine_parameter_functor(), -10.0, 10.0, + nullptr, scale, offset); + + std::vector inputs{scale, offset}; + std::vector gradient; + root.grad(inputs, gradient); + + EXPECT_NEAR(2.0 / 3.0, root.val(), 1e-10); + EXPECT_NEAR(-2.0 / 9.0, gradient[0], 1e-10); + EXPECT_NEAR(1.0 / 3.0, gradient[1], 1e-10); + stan::math::recover_memory(); +} + +TEST(MathRevFunctorSolveItp, ImplicitDerivativeVectorParameter) { + using stan::math::var; + Eigen::Matrix theta(3); + theta << 2.0, 1.0, 5.0; + var root = stan::math::solve_itp(vector_parameter_functor(), -10.0, 10.0, + nullptr, theta); + + std::vector inputs{theta(0), theta(1), theta(2)}; + std::vector gradient; + root.grad(inputs, gradient); + + EXPECT_NEAR(2.0, root.val(), 1e-10); + EXPECT_NEAR(-1.0, gradient[0], 1e-10); + EXPECT_NEAR(-0.5, gradient[1], 1e-10); + EXPECT_NEAR(0.5, gradient[2], 1e-10); + stan::math::recover_memory(); +} + +TEST(MathRevFunctorSolveItp, RejectsSingularImplicitDerivative) { + using stan::math::var; + var theta = 0.0; + EXPECT_THROW(stan::math::solve_itp(singular_parameter_functor(), -1.0, 0.0, + nullptr, theta), + std::domain_error); + stan::math::recover_memory(); +}