diff --git a/stan/math/fwd/fun/Eigen_NumTraits.hpp b/stan/math/fwd/fun/Eigen_NumTraits.hpp index 9e627e46d1e..d4577f20d4c 100644 --- a/stan/math/fwd/fun/Eigen_NumTraits.hpp +++ b/stan/math/fwd/fun/Eigen_NumTraits.hpp @@ -253,4 +253,5 @@ struct ternary_evaluator< } // namespace internal } // namespace Eigen + #endif diff --git a/stan/math/prim/fun.hpp b/stan/math/prim/fun.hpp index 8ef4204c81b..0a384317105 100644 --- a/stan/math/prim/fun.hpp +++ b/stan/math/prim/fun.hpp @@ -223,8 +223,10 @@ #include #include #include +#include #include #include +#include #include #include #include @@ -281,6 +283,7 @@ #include #include #include +#include #include #include #include diff --git a/stan/math/prim/fun/multiply_add.hpp b/stan/math/prim/fun/multiply_add.hpp new file mode 100644 index 00000000000..534f19f4d82 --- /dev/null +++ b/stan/math/prim/fun/multiply_add.hpp @@ -0,0 +1,46 @@ +#ifndef STAN_MATH_PRIM_FUN_MULTIPLY_ADD_HPP +#define STAN_MATH_PRIM_FUN_MULTIPLY_ADD_HPP + +#include +#include +#include + +namespace stan { +namespace math { + +/** + * Return the product of two matrices plus a third, `A * B + C`. The + * argument order follows `fma(x, y, z) = x * y + z`. + * + * @tparam Mat1 type of the first matrix or expression + * @tparam Mat2 type of the second matrix or expression + * @tparam Mat3 type of the matrix or expression added to the product + * @param A first matrix + * @param B second matrix + * @param C matrix added to the product + * @return `A * B + C` + * @throw std::invalid_argument if the number of columns of `A` does not + * match the number of rows of `B`, or `C` does not have the dimensions of + * `A * B` + */ +template * = nullptr, + require_all_st_arithmetic* = nullptr> +inline auto multiply_add(Mat1&& A, Mat2&& B, Mat3&& C) { + check_multiplicable("multiply_add", "A", A, "B", B); + check_size_match("multiply_add", "Rows of A", A.rows(), "rows of C", + C.rows()); + check_size_match("multiply_add", "Columns of B", B.cols(), "columns of C", + C.cols()); + return make_holder( + [](auto&& A_, auto&& B_, auto&& C_) { + return std::forward(A_) * std::forward(B_) + + std::forward(C_); + }, + std::forward(A), std::forward(B), std::forward(C)); +} + +} // namespace math +} // namespace stan + +#endif diff --git a/stan/math/prim/fun/multiply_sub.hpp b/stan/math/prim/fun/multiply_sub.hpp new file mode 100644 index 00000000000..a1abd2231f1 --- /dev/null +++ b/stan/math/prim/fun/multiply_sub.hpp @@ -0,0 +1,48 @@ +#ifndef STAN_MATH_PRIM_FUN_MULTIPLY_SUB_HPP +#define STAN_MATH_PRIM_FUN_MULTIPLY_SUB_HPP + +#include +#include +#include + +namespace stan { +namespace math { + +/** + * Return a matrix minus the product of two matrices, `C - A * B`. The + * argument order follows `multiply_add(A, B, C) = A * B + C`. + * + * @tparam Mat1 type of the first matrix or expression + * @tparam Mat2 type of the second matrix or expression + * @tparam Mat3 type of the matrix or expression the product is subtracted + * from + * @param A first matrix + * @param B second matrix + * @param C matrix the product is subtracted from + * @return `C - A * B` + * @throw std::invalid_argument if the number of columns of `A` does not + * match the number of rows of `B`, or `C` does not have the dimensions of + * `A * B` + */ +template * = nullptr, + require_all_st_arithmetic* = nullptr> +inline auto multiply_sub(Mat1&& A, Mat2&& B, Mat3&& C) { + check_multiplicable("multiply_sub", "A", A, "B", B); + check_size_match("multiply_sub", "Rows of A", A.rows(), "rows of C", + C.rows()); + check_size_match("multiply_sub", "Columns of B", B.cols(), "columns of C", + C.cols()); + return make_holder( + [](auto&& A_, auto&& B_, auto&& C_) { + return std::forward(C_) + - std::forward(A_) + * std::forward(B_); + }, + std::forward(A), std::forward(B), std::forward(C)); +} + +} // namespace math +} // namespace stan + +#endif diff --git a/stan/math/prim/fun/scale_multiply.hpp b/stan/math/prim/fun/scale_multiply.hpp new file mode 100644 index 00000000000..6c080597d20 --- /dev/null +++ b/stan/math/prim/fun/scale_multiply.hpp @@ -0,0 +1,42 @@ +#ifndef STAN_MATH_PRIM_FUN_SCALE_MULTIPLY_HPP +#define STAN_MATH_PRIM_FUN_SCALE_MULTIPLY_HPP + +#include +#include +#include + +namespace stan { +namespace math { + +/** + * Return the product of a scalar and two matrices, `s * A * B`. + * + * @tparam Scal type of the scalar + * @tparam Mat1 type of the first matrix or expression + * @tparam Mat2 type of the second matrix or expression + * @param s scalar + * @param A first matrix + * @param B second matrix + * @return `s * A * B` + * @throw std::invalid_argument if the number of columns of `A` does not + * match the number of rows of `B` + */ +template * = nullptr, + require_all_eigen_t* = nullptr, + require_all_st_arithmetic* = nullptr> +inline auto scale_multiply(Scal&& s, Mat1&& A, Mat2&& B) { + check_multiplicable("scale_multiply", "A", A, "B", B); + return make_holder( + [s](auto&& A_, auto&& B_) { + return s + * (std::forward(A_) + * std::forward(B_)); + }, + std::forward(A), std::forward(B)); +} + +} // namespace math +} // namespace stan + +#endif diff --git a/stan/math/prim/meta/is_eigen.hpp b/stan/math/prim/meta/is_eigen.hpp index ef776763c64..81f7ae88d8c 100644 --- a/stan/math/prim/meta/is_eigen.hpp +++ b/stan/math/prim/meta/is_eigen.hpp @@ -261,7 +261,11 @@ using require_all_not_eigen_st = require_all_not_t< */ template struct is_eigen_array - : bool_constant::value> {}; + : bool_constant>::value> {}; + +template +inline constexpr bool is_eigen_array_v = is_eigen_array::value; /*! \ingroup require_eigens_types */ /*! \defgroup eigen_array_types eigen_array */ diff --git a/stan/math/rev/core/Eigen_NumTraits.hpp b/stan/math/rev/core/Eigen_NumTraits.hpp index ed5eebad4af..cace585b582 100644 --- a/stan/math/rev/core/Eigen_NumTraits.hpp +++ b/stan/math/rev/core/Eigen_NumTraits.hpp @@ -5,6 +5,7 @@ #include #include #include +#include #include #include #include @@ -504,4 +505,5 @@ struct general_matrix_matrix_product #include #include +#include #include #include +#include #include #include #include @@ -149,6 +151,7 @@ #include #include #include +#include #include #include #include diff --git a/stan/math/rev/fun/multiply_add.hpp b/stan/math/rev/fun/multiply_add.hpp new file mode 100644 index 00000000000..02e68ba5a25 --- /dev/null +++ b/stan/math/rev/fun/multiply_add.hpp @@ -0,0 +1,68 @@ +#ifndef STAN_MATH_REV_FUN_MULTIPLY_ADD_HPP +#define STAN_MATH_REV_FUN_MULTIPLY_ADD_HPP + +#include +#include +#include +#include +#include +#include + +namespace stan { +namespace math { + +/** + * Return the product of two matrices plus a third, `A * B + C`, where at + * least one argument has `var` scalars. The whole expression is one node on + * the reverse-mode stack. The argument order follows `fma(x, y, z) = x * y + + * z`. + * + * @tparam Mat1 type of the first matrix or expression + * @tparam Mat2 type of the second matrix or expression + * @tparam Mat3 type of the matrix or expression added to the product + * @param A first matrix + * @param B second matrix + * @param C matrix added to the product + * @return `A * B + C` + * @throw std::invalid_argument if the number of columns of `A` does not + * match the number of rows of `B`, or `C` does not have the dimensions of + * `A * B` + */ +template * = nullptr, + require_any_st_var* = nullptr, + require_all_not_complex_t, value_type_t, + value_type_t>* = nullptr> +inline auto multiply_add(Mat1&& A, Mat2&& B, Mat3&& C) { + check_multiplicable("multiply_add", "A", A, "B", B); + check_size_match("multiply_add", "Rows of A", A.rows(), "rows of C", + C.rows()); + check_size_match("multiply_add", "Columns of B", B.cols(), "columns of C", + C.cols()); + arena_t arena_A(std::forward(A)); + arena_t arena_B(std::forward(B)); + arena_t arena_C(std::forward(C)); + using return_t + = return_var_matrix_t; + arena_t res + = value_of(arena_A) * value_of(arena_B) + value_of(arena_C); + reverse_pass_callback([arena_A, arena_B, arena_C, res]() mutable { + // Eigen's matrix product cannot read the adjoint view directly + const auto res_adj = res.adj().eval(); + if constexpr (is_autodiff_v) { + arena_A.adj() += res_adj * value_of(arena_B).transpose(); + } + if constexpr (is_autodiff_v) { + arena_B.adj() += value_of(arena_A).transpose() * res_adj; + } + if constexpr (is_autodiff_v) { + arena_C.adj() += res_adj; + } + }); + return res; +} + +} // namespace math +} // namespace stan + +#endif diff --git a/stan/math/rev/fun/multiply_sub.hpp b/stan/math/rev/fun/multiply_sub.hpp new file mode 100644 index 00000000000..e505a32ceef --- /dev/null +++ b/stan/math/rev/fun/multiply_sub.hpp @@ -0,0 +1,69 @@ +#ifndef STAN_MATH_REV_FUN_MULTIPLY_SUB_HPP +#define STAN_MATH_REV_FUN_MULTIPLY_SUB_HPP + +#include +#include +#include +#include +#include +#include + +namespace stan { +namespace math { + +/** + * Return a matrix minus the product of two matrices, `C - A * B`, where at + * least one argument has `var` scalars. The whole expression is one node on + * the reverse-mode stack. The argument order follows + * `multiply_add(A, B, C) = A * B + C`. + * + * @tparam Mat1 type of the first matrix or expression + * @tparam Mat2 type of the second matrix or expression + * @tparam Mat3 type of the matrix or expression the product is subtracted + * from + * @param A first matrix + * @param B second matrix + * @param C matrix the product is subtracted from + * @return `C - A * B` + * @throw std::invalid_argument if the number of columns of `A` does not + * match the number of rows of `B`, or `C` does not have the dimensions of + * `A * B` + */ +template * = nullptr, + require_any_st_var* = nullptr, + require_all_not_complex_t, value_type_t, + value_type_t>* = nullptr> +inline auto multiply_sub(Mat1&& A, Mat2&& B, Mat3&& C) { + check_multiplicable("multiply_sub", "A", A, "B", B); + check_size_match("multiply_sub", "Rows of A", A.rows(), "rows of C", + C.rows()); + check_size_match("multiply_sub", "Columns of B", B.cols(), "columns of C", + C.cols()); + arena_t arena_A(std::forward(A)); + arena_t arena_B(std::forward(B)); + arena_t arena_C(std::forward(C)); + using return_t + = return_var_matrix_t; + arena_t res + = value_of(arena_C) - value_of(arena_A) * value_of(arena_B); + reverse_pass_callback([arena_A, arena_B, arena_C, res]() mutable { + // Eigen's matrix product cannot read the adjoint view directly + const auto res_adj = res.adj().eval(); + if constexpr (is_autodiff_v) { + arena_A.adj() -= res_adj * value_of(arena_B).transpose(); + } + if constexpr (is_autodiff_v) { + arena_B.adj() -= value_of(arena_A).transpose() * res_adj; + } + if constexpr (is_autodiff_v) { + arena_C.adj() += res_adj; + } + }); + return res; +} + +} // namespace math +} // namespace stan + +#endif diff --git a/stan/math/rev/fun/scale_multiply.hpp b/stan/math/rev/fun/scale_multiply.hpp new file mode 100644 index 00000000000..10629fc9cbd --- /dev/null +++ b/stan/math/rev/fun/scale_multiply.hpp @@ -0,0 +1,62 @@ +#ifndef STAN_MATH_REV_FUN_SCALE_MULTIPLY_HPP +#define STAN_MATH_REV_FUN_SCALE_MULTIPLY_HPP + +#include +#include +#include +#include +#include +#include + +namespace stan { +namespace math { + +/** + * Return the product of a scalar and two matrices, `s * A * B`, where at + * least one argument has `var` scalars. The whole product is one node on + * the reverse-mode stack. + * + * @tparam Scal type of the scalar + * @tparam Mat1 type of the first matrix or expression + * @tparam Mat2 type of the second matrix or expression + * @param s scalar + * @param A first matrix + * @param B second matrix + * @return `s * A * B` + * @throw std::invalid_argument if the number of columns of `A` does not + * match the number of rows of `B` + */ +template * = nullptr, + require_all_eigen_t* = nullptr, + require_any_st_var* = nullptr, + require_all_not_complex_t, + value_type_t>* = nullptr> +inline auto scale_multiply(Scal&& s, Mat1&& A, Mat2&& B) { + check_multiplicable("scale_multiply", "A", A, "B", B); + arena_t arena_A(std::forward(A)); + arena_t arena_B(std::forward(B)); + using return_t = return_var_matrix_t; + arena_t res = value_of(s) * (value_of(arena_A) * value_of(arena_B)); + reverse_pass_callback([s, arena_A, arena_B, res]() mutable { + // Eigen's matrix product cannot read the adjoint view directly + const auto res_adj = res.adj().eval(); + if constexpr (is_autodiff_v) { + s.adj() += (value_of(arena_A).transpose() * res_adj) + .cwiseProduct(value_of(arena_B)) + .sum(); + } + if constexpr (is_autodiff_v) { + arena_A.adj() += value_of(s) * res_adj * value_of(arena_B).transpose(); + } + if constexpr (is_autodiff_v) { + arena_B.adj() += value_of(s) * value_of(arena_A).transpose() * res_adj; + } + }); + return res; +} + +} // namespace math +} // namespace stan + +#endif diff --git a/stan/math/rev/meta.hpp b/stan/math/rev/meta.hpp index 892080b5343..dea174f7c2f 100644 --- a/stan/math/rev/meta.hpp +++ b/stan/math/rev/meta.hpp @@ -4,6 +4,7 @@ #include #include #include +#include #include #include #include diff --git a/stan/math/rev/meta/eigen_product_evaluator_hooks.hpp b/stan/math/rev/meta/eigen_product_evaluator_hooks.hpp new file mode 100644 index 00000000000..c040a75c7c8 --- /dev/null +++ b/stan/math/rev/meta/eigen_product_evaluator_hooks.hpp @@ -0,0 +1,468 @@ +#ifndef STAN_MATH_REV_META_EIGEN_PRODUCT_EVALUATOR_HOOKS_HPP +#define STAN_MATH_REV_META_EIGEN_PRODUCT_EVALUATOR_HOOKS_HPP + +/** + * @file + * Specializations of Eigen's evaluators that compute products of matrices + * with autodiff scalars with Stan's `multiply`, `elt_multiply`, + * `scale_multiply`, `multiply_add` and `multiply_sub`, so that `A * B` + * records one node on the reverse-mode stack instead of one per scalar + * operation. + * + * Each specialization must be more specialized than Eigen's under partial + * ordering. A condition (`enable_if_t`) can only go in an argument that + * Eigen's specialization leaves free; where Eigen's has a computed default + * or its own condition, ours must repeat it exactly or the two are + * ambiguous. When upgrading Eigen, compare the specializations in + * Eigen/src/Core/ProductEvaluators.h and CoreEvaluators.h with these. + * + * This file is included by stan/math/rev/meta.hpp and + * stan/math/rev/core/Eigen_NumTraits.hpp, so every translation unit that + * uses `var` sees it. The functions the specializations call are found by + * argument-dependent lookup when a product is evaluated, so a translation + * unit that evaluates a `var` product must include stan/math/rev.hpp. + */ + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +namespace stan { +namespace math { +namespace internal { + +/** + * True if either operand of an Eigen product has `var` scalars and neither + * has complex scalars, so the product is computed by Stan's functions. + * Products with a complex operand are left to Eigen: Stan's `multiply` for + * them evaluates an Eigen `lazyProduct`, which would come back here. + * + * @tparam Lhs left operand expression type + * @tparam Rhs right operand expression type + */ +template +inline constexpr bool is_stan_ad_product_v + = is_var>::value; + +template +struct is_scalar_times_matrix_impl : std::false_type {}; +template +struct is_scalar_times_matrix_impl, + const Eigen::CwiseNullaryOp, Plain>, + const Mat>> : std::true_type {}; + +/** + * True for the expression Eigen builds for a scalar times a matrix, `s * X`: + * a `CwiseBinaryOp` of `scalar_product_op` whose left operand is a + * `CwiseNullaryOp` holding `s` + * (`xpr.lhs().functor().m_other`) and whose right operand is `X` (`xpr.rhs()`). + * + * @tparam T an Eigen expression type, with any reference or `const` + */ +template +struct is_scalar_times_matrix : is_scalar_times_matrix_impl> {}; + +template +struct is_matrix_times_scalar_impl : std::false_type {}; +template +struct is_matrix_times_scalar_impl< + Eigen::CwiseBinaryOp, const Mat, + const Eigen::CwiseNullaryOp< + Eigen::internal::scalar_constant_op, Plain>>> + : std::true_type {}; + +/** + * True for the expression Eigen builds for a matrix times a scalar, `X * s`: + * a `CwiseBinaryOp` of `scalar_product_op` whose left operand is `X` + * (`xpr.lhs()`) and whose right operand is a + * `CwiseNullaryOp` holding `s` + * (`xpr.rhs().functor().m_other`). + * + * @tparam T an Eigen expression type, with any reference or `const` + */ +template +struct is_matrix_times_scalar : is_matrix_times_scalar_impl> {}; + +template +struct is_eigen_product_impl : std::false_type {}; +template +struct is_eigen_product_impl> + : std::true_type {}; + +/** + * True for an Eigen matrix product expression. + * + * @tparam T an Eigen expression type, with any reference or `const` + */ +template +struct is_eigen_product : is_eigen_product_impl> {}; + +template +struct is_eigen_cwise_product_impl : std::false_type {}; +template +struct is_eigen_cwise_product_impl, Lhs, Rhs>> + : std::true_type {}; + +/** + * True for an Eigen coefficient-wise product expression: `A.cwiseProduct(B)`, + * `a * b` of arrays, and `s * X` / `X * s`, which Eigen builds as a + * coefficient-wise product with a constant expression. + * + * @tparam T an Eigen expression type, with any reference or `const` + */ +template +struct is_eigen_cwise_product : is_eigen_cwise_product_impl> {}; + +/*! \brief Require type satisfies @ref is_eigen_product */ +/*! @tparam T the type to check */ +template +using require_eigen_product_t = require_t>; + +/*! \brief Require type satisfies @ref is_eigen_cwise_product */ +/*! @tparam T the type to check */ +template +using require_eigen_cwise_product_t = require_t>; + +/** + * Empty base of every Stan specialization of an Eigen evaluator or + * assignment below, so tests can check which specialization Eigen picks for + * an expression (`std::is_base_of>`). + */ +struct eigen_ad_hook {}; + +/** + * `var` if Eigen's assignment of the product of `Lhs` and `Rhs` is computed + * by Stan's functions (both operands are dense), otherwise no type. Used as + * both scalars of the assignment functor in the `Assignment` + * specializations below. + * + * @tparam Lhs left operand expression type + * @tparam Rhs right operand expression type + */ +template +using var_if_dense_product_t = std::enable_if_t< + std::is_same::Shape, + Eigen::DenseShape>::value + && std::is_same::Shape, + Eigen::DenseShape>::value, + var>; + +/** + * Return `x`, or an arena copy of its vars if `x` is an `Eigen::Map` of vars. + * + * Stan's functions assume an `Eigen::Map` points at arena memory and keep it + * in their reverse-pass callbacks without copying (see + * `arena_matrix(const Map&)`), so the callback reads whatever vars the mapped + * memory holds during the reverse pass. Eigen's own algorithms map + * workspaces that they reuse and free, e.g. the Householder reflections in + * `HouseholderQR`. Copying the vars (not their varis) keeps the ones the + * forward pass used. + * + * @tparam T type of the operand + * @param x an operand passed from an Eigen expression to a Stan function + */ +template +inline decltype(auto) copy_var_map(T&& x) { + if constexpr (is_eigen_contiguous_map::value + && is_var>::value) { + arena_matrix> vars(x.rows(), x.cols()); + vars.deep_copy(x); + return vars; + } else { + return std::forward(x); + } +} + +/** + * Return the product of the operands of an Eigen matrix product. When an + * operand is Eigen's `s * A` or `A * s` (Eigen rewrites `s * (A * B)` as + * `(s * A) * B`) this is `scale_multiply(s, A, B)`, otherwise + * `multiply(lhs, rhs)`. + * + * @tparam Lhs left operand expression type + * @tparam Rhs right operand expression type + * @param lhs left operand + * @param rhs right operand + */ +template +inline auto multiply_operands(Lhs&& lhs, Rhs&& rhs) { + if constexpr (is_scalar_times_matrix::value) { + // (s * A) * B + return scale_multiply(std::forward(lhs).lhs().functor().m_other, + copy_var_map(std::forward(lhs).rhs()), + copy_var_map(std::forward(rhs))); + } else if constexpr (is_matrix_times_scalar::value) { + // (A * s) * B + return scale_multiply(std::forward(lhs).rhs().functor().m_other, + copy_var_map(std::forward(lhs).lhs()), + copy_var_map(std::forward(rhs))); + } else if constexpr (is_scalar_times_matrix::value) { + // A * (s * B) + return scale_multiply(std::forward(rhs).lhs().functor().m_other, + copy_var_map(std::forward(lhs)), + copy_var_map(std::forward(rhs).rhs())); + } else if constexpr (is_matrix_times_scalar::value) { + // A * (B * s) + return scale_multiply(std::forward(rhs).rhs().functor().m_other, + copy_var_map(std::forward(lhs)), + copy_var_map(std::forward(rhs).lhs())); + } else { + return multiply(copy_var_map(std::forward(lhs)), + copy_var_map(std::forward(rhs))); + } +} + +/** + * Return the matrix product `xpr` computed by `multiply_operands`. + * + * @param xpr an Eigen matrix product + */ +template * = nullptr> +inline auto evaluate_product(XprType&& xpr) { + return multiply_operands(std::forward(xpr).lhs(), + std::forward(xpr).rhs()); +} + +/** + * Return the coefficient-wise product `xpr`: `scale_multiply(s, A, B)` for + * `s * (A * B)` or `(A * B) * s`, `multiply(s, X)` for `s * X` or `X * s`, + * and `elt_multiply(lhs, rhs)` otherwise. Stan's functions take matrices, so + * array operands are viewed as matrices (`.matrix().derived()` is the + * operand itself for a matrix). + * + * @param xpr an Eigen coefficient-wise product + */ +template * = nullptr> +inline auto evaluate_product(XprType&& xpr) { + // the operand types, without a reference or `const` + using Xpr = std::decay_t; + if constexpr (is_scalar_times_matrix::value) { + if constexpr (is_eigen_product::value) { + // s * (A * B) + return scale_multiply( + std::forward(xpr).lhs().functor().m_other, + copy_var_map(std::forward(xpr).rhs().lhs()), + copy_var_map(std::forward(xpr).rhs().rhs())); + } else { + // s * X + return multiply( + std::forward(xpr).lhs().functor().m_other, + copy_var_map(std::forward(xpr).rhs().matrix().derived())); + } + } else if constexpr (is_matrix_times_scalar::value) { + if constexpr (is_eigen_product::value) { + // (A * B) * s + return scale_multiply( + std::forward(xpr).rhs().functor().m_other, + copy_var_map(std::forward(xpr).lhs().lhs()), + copy_var_map(std::forward(xpr).lhs().rhs())); + } else { + // X * s + return multiply( + std::forward(xpr).rhs().functor().m_other, + copy_var_map(std::forward(xpr).lhs().matrix().derived())); + } + } else { + return elt_multiply( + copy_var_map(std::forward(xpr).lhs().matrix().derived()), + copy_var_map(std::forward(xpr).rhs().matrix().derived())); + } +} + +/** + * Body of the `product_evaluator` and `binary_evaluator` specializations + * below. Like Eigen's product evaluator it holds the product, computed by + * `evaluate_product`, and evaluates to it. + * + * @tparam XprType an Eigen matrix or coefficient-wise product + */ +template +struct ad_product_evaluator + : public Eigen::internal::evaluator, + eigen_ad_hook { + using PlainObject = typename XprType::PlainObject; + using Base = Eigen::internal::evaluator; + enum { Flags = Base::Flags | Eigen::EvalBeforeNestingBit }; + template , + std::decay_t>>* = nullptr> + explicit ad_product_evaluator(XprTypeT&& xpr) { + if constexpr (is_stan_scalar_v) { + // a row vector times a column vector + m_result.resize(1, 1); + m_result.coeffRef(0, 0) = evaluate_product(std::forward(xpr)); + } else if constexpr (is_eigen_array_v) { + m_result = evaluate_product(std::forward(xpr)).array(); + } else { + m_result = evaluate_product(std::forward(xpr)); + } + // Base points into m_result's storage, so it is constructed only once + // m_result is filled: assigning a returned matrix can move a new buffer + // into m_result. + Eigen::internal::construct_at(this, m_result); + } + + protected: + PlainObject m_result; +}; + +} // namespace internal +} // namespace math +} // namespace stan + +namespace Eigen { +namespace internal { + +/** + * Dense-by-dense products where an operand has `var` scalars, computed by + * `ad_product_evaluator` with `multiply_operands` (`multiply` or + * `scale_multiply`). + * + * The `var` condition sits in the left shape argument: Eigen's generic + * `product_evaluator` takes any type there, so this specialization is more + * specialized than it and does not exist for other scalars. (Putting the + * condition in the scalar arguments or in the type of `Options` makes the + * two specializations ambiguous, because Eigen's generic one computes the + * scalars and takes an `int` there.) + * + * Lazy products (`lazyProduct()`, and the `LazyCoeffBasedProductMode` + * evaluator that derives from Eigen's lazy one) are left to Eigen: its lazy + * evaluators fix both shapes, so this specialization would be ambiguous with + * them. + */ +template +struct product_evaluator< + Product, ProductTag, + std::enable_if_t && Options != LazyProduct + && ProductTag != LazyCoeffBasedProductMode, + DenseShape>, + DenseShape> + : stan::math::internal::ad_product_evaluator> { + using XprType = Product; + template * = nullptr> + explicit product_evaluator(XprTypeT&& xpr) + : stan::math::internal::ad_product_evaluator( + std::forward(xpr)) {} +}; + +/** + * Coefficient-wise products `lhs .* rhs` where an operand has `var` scalars + * and neither is complex, computed by `ad_product_evaluator`. + * + * The condition sits in the functor argument of `CwiseBinaryOp`, which + * Eigen's specialization leaves free, so this one is more specialized. The + * functor is written as Eigen builds it, `scalar_product_op` of the two + * operands' scalars, so `Lhs` and `Rhs` are deduced from the operands. + */ +template +struct binary_evaluator< + CwiseBinaryOp< + std::enable_if_t, + scalar_product_op::Scalar, + typename traits::Scalar>>, + Lhs, Rhs>, + IndexBased, IndexBased> + : stan::math::internal::ad_product_evaluator< + CwiseBinaryOp::Scalar, + typename traits::Scalar>, + Lhs, Rhs>> { + using XprType = CwiseBinaryOp::Scalar, + typename traits::Scalar>, + Lhs, Rhs>; + explicit binary_evaluator(const XprType& xpr) + : stan::math::internal::ad_product_evaluator(xpr) {} +}; + +/** + * Eigen assigns a product to a destination (`dst = A * B`, and + * `dst.noalias() += A * B` / `-=`) with its own `Assignment` + * specializations, which compute the product without `product_evaluator`. + * Eigen has one specialization per assignment functor, so these are three + * as well. The condition (`var` product of dense operands) sits in both + * scalars of the functor, which Eigen's specializations leave free, so these + * are more specialized; other products run Eigen's own specialization. The + * last argument must repeat Eigen's own condition token for token (a + * different condition makes the two ambiguous). + */ +template +struct Assignment< + DstXprType, Product, + assign_op, + stan::math::internal::var_if_dense_product_t>, + Dense2Dense, + std::enable_if_t<(Options == DefaultProduct + || Options == AliasFreeProduct)>> + : stan::math::internal::eigen_ad_hook { + template * = nullptr> + static void run(DstXprType& dst, SrcXpr&& src, Func&& /* func */) { + using stan::math::internal::multiply_operands; + if constexpr (stan::is_stan_scalar::value) { + // a row vector times a column vector + if (dst.rows() != 1 || dst.cols() != 1) { + dst.resize(1, 1); + } + dst.coeffRef(0, 0) = multiply_operands(std::forward(src).lhs(), + std::forward(src).rhs()); + } else { + dst = multiply_operands(std::forward(src).lhs(), + std::forward(src).rhs()); + } + } +}; + +template +struct Assignment< + DstXprType, Product, + add_assign_op, + stan::math::internal::var_if_dense_product_t>, + Dense2Dense, + std::enable_if_t<(Options == DefaultProduct + || Options == AliasFreeProduct)>> + : stan::math::internal::eigen_ad_hook { + template * = nullptr> + static void run(DstXprType& dst, SrcType&& src, Func&& /* func */) { + using stan::math::internal::copy_var_map; + dst = multiply_add(copy_var_map(std::forward(src).lhs()), + copy_var_map(std::forward(src).rhs()), + copy_var_map(dst)); + } +}; + +template +struct Assignment< + DstXprType, Product, + sub_assign_op, + stan::math::internal::var_if_dense_product_t>, + Dense2Dense, + std::enable_if_t<(Options == DefaultProduct + || Options == AliasFreeProduct)>> + : stan::math::internal::eigen_ad_hook { + template * = nullptr> + static void run(DstXprType& dst, SrcType&& src, Func&& /* func */) { + using stan::math::internal::copy_var_map; + dst = multiply_sub(copy_var_map(std::forward(src).lhs()), + copy_var_map(std::forward(src).rhs()), + copy_var_map(dst)); + } +}; + +} // namespace internal +} // namespace Eigen + +#endif diff --git a/test/unit/math/mix/core/eigen_product_hook_kernels_tape_test.cpp b/test/unit/math/mix/core/eigen_product_hook_kernels_tape_test.cpp new file mode 100644 index 00000000000..6c3d1f0842e --- /dev/null +++ b/test/unit/math/mix/core/eigen_product_hook_kernels_tape_test.cpp @@ -0,0 +1,83 @@ +#include +#include + +// Eigen's own kernels (triangular products and solves, Cholesky, LU) form +// coefficient-wise and scalar products internally, which Stan's product +// hooks evaluate. These bounds are the number of reverse-pass (chain) varis +// each kernel recorded before the hooks existed, for n = 8; the hooks must +// not increase them. + +namespace { +using stan::math::fvar; +using stan::math::var; + +template +T kernel_scalar(double x) { + return T(x); +} +template <> +fvar kernel_scalar>(double x) { + return fvar(var(x), var(0.5)); +} + +struct kernel_chain_counts { + std::size_t tri_solve; + std::size_t tri_mul; + std::size_t llt; + std::size_t lu; +}; + +template +kernel_chain_counts measure_kernels() { + const int n = 8; + Eigen::Matrix a(n, n); + Eigen::Matrix b(n); + for (int j = 0; j < n; ++j) { + b(j) = kernel_scalar(0.1 * j + 0.3); + for (int i = 0; i < n; ++i) { + a(i, j) = kernel_scalar((i == j ? n : 0) + 0.01 * (i + 1) * (j + 2)); + } + } + const auto& chain_stack = stan::math::ChainableStack::instance_->var_stack_; + kernel_chain_counts counts; + { + const std::size_t start = chain_stack.size(); + Eigen::Matrix x + = a.template triangularView().solve(b); + counts.tri_solve = chain_stack.size() - start; + } + { + const std::size_t start = chain_stack.size(); + Eigen::Matrix x = a.template triangularView() * b; + counts.tri_mul = chain_stack.size() - start; + } + { + const std::size_t start = chain_stack.size(); + Eigen::Matrix x = a.llt().solve(b); + counts.llt = chain_stack.size() - start; + } + { + const std::size_t start = chain_stack.size(); + Eigen::Matrix x = a.partialPivLu().solve(b); + counts.lu = chain_stack.size() - start; + } + stan::math::recover_memory(); + return counts; +} +} // namespace + +TEST(mathMixEigenProductHook, varKernelChainCountsDoNotGrow) { + const auto counts = measure_kernels(); + EXPECT_LE(counts.tri_solve, 64); + EXPECT_LE(counts.tri_mul, 82); + EXPECT_LE(counts.llt, 340); + EXPECT_LE(counts.lu, 484); +} + +TEST(mathMixEigenProductHook, fvarVarKernelChainCountsDoNotGrow) { + const auto counts = measure_kernels>(); + EXPECT_LE(counts.tri_solve, 216); + EXPECT_LE(counts.tri_mul, 256); + EXPECT_LE(counts.llt, 1414); + EXPECT_LE(counts.lu, 1504); +} diff --git a/test/unit/math/mix/fun/multiply_add_test.cpp b/test/unit/math/mix/fun/multiply_add_test.cpp new file mode 100644 index 00000000000..81212aeeb69 --- /dev/null +++ b/test/unit/math/mix/fun/multiply_add_test.cpp @@ -0,0 +1,21 @@ +#include + +TEST(mathMixMatFun, multiplyAdd) { + auto f = [](const auto& a, const auto& b, const auto& c) { + return stan::math::multiply_add(a, b, c); + }; + Eigen::MatrixXd a(3, 2); + a << 1, -2, 0.5, 3, -1.5, 2; + Eigen::MatrixXd b(2, 2); + b << 0.25, -1, 2, 0.5; + Eigen::MatrixXd c(3, 2); + c << 0.1, 0.2, -0.3, 0.4, 0.5, -0.6; + Eigen::VectorXd v(2); + v << -0.5, 1.5; + Eigen::VectorXd cv(3); + cv << 1, 2, 3; + // reverse mode only: multiply_add has no fvar overload yet + constexpr bool reverse_only = true; + stan::test::expect_ad(f, a, b, c); + stan::test::expect_ad(f, a, v, cv); +} diff --git a/test/unit/math/mix/fun/multiply_operator_cwise_test.cpp b/test/unit/math/mix/fun/multiply_operator_cwise_test.cpp new file mode 100644 index 00000000000..45e850871a4 --- /dev/null +++ b/test/unit/math/mix/fun/multiply_operator_cwise_test.cpp @@ -0,0 +1,56 @@ +#include +#include + +// Values and derivatives of coefficient-wise products and products with a +// scalar of autodiff matrices and arrays. Which Stan function computes each +// form is tested in test/unit/math/rev/core/eigen_product_hook_test.cpp. + +using stan::test::multiply_operator::test_matrix; + +TEST(mathMixMultiplyOperator, arrayTimesArray) { + auto f = [](const auto& a, const auto& b) { + return (a.array() * b.array()).matrix().eval(); + }; + Eigen::MatrixXd a = test_matrix(3, 2, 0.1); + Eigen::MatrixXd b = test_matrix(3, 2, -0.4); + stan::test::expect_ad(f, a, b); +} + +TEST(mathMixMultiplyOperator, cwiseProduct) { + auto f + = [](const auto& a, const auto& b) { return a.cwiseProduct(b).eval(); }; + Eigen::MatrixXd a = test_matrix(3, 2, 0.1); + Eigen::MatrixXd b = test_matrix(3, 2, -0.4); + stan::test::expect_ad(f, a, b); +} + +TEST(mathMixMultiplyOperator, cwiseNestedAndPartialRead) { + auto f_nested = [](const auto& a, const auto& b) { + return ((a.array() * b.array()).exp() * a.array()).matrix().eval(); + }; + auto f_head = [](const auto& a, const auto& b) { + return (a.array() * b.array()).matrix().col(0).head(2).eval(); + }; + Eigen::MatrixXd a = test_matrix(3, 2, 0.1); + Eigen::MatrixXd b = test_matrix(3, 2, -0.4); + stan::test::expect_ad(f_nested, a, b); + stan::test::expect_ad(f_head, a, b); +} + +TEST(mathMixMultiplyOperator, scalarTimesMatrix) { + auto left = [](const auto& s, const auto& a) { return (s * a).eval(); }; + auto right = [](const auto& s, const auto& a) { return (a * s).eval(); }; + auto array = [](const auto& s, const auto& a) { + return (a.array() * s).matrix().eval(); + }; + double s = 1.3; + Eigen::MatrixXd a = test_matrix(2, 3, 0.2); + stan::test::expect_ad(left, s, a); + stan::test::expect_ad(right, s, a); + stan::test::expect_ad(array, s, a); +} + +TEST(mathMixMultiplyOperator, matrixTimesInt) { + auto f = [](const auto& a) { return (a * 3).eval(); }; + stan::test::expect_ad(f, test_matrix(2, 3, 0.2)); +} diff --git a/test/unit/math/mix/fun/multiply_operator_matmul_test.cpp b/test/unit/math/mix/fun/multiply_operator_matmul_test.cpp new file mode 100644 index 00000000000..9f9b2ee6a6a --- /dev/null +++ b/test/unit/math/mix/fun/multiply_operator_matmul_test.cpp @@ -0,0 +1,99 @@ +#include +#include + +// Values and derivatives of Eigen `operator*` matrix products of autodiff +// matrices. Which Stan function computes each form is tested in +// test/unit/math/rev/core/eigen_product_hook_test.cpp. + +using stan::test::multiply_operator::test_matrix; + +namespace { +// Eigen cannot multiply fvar and double matrices until fvar products are +// routed to Stan's functions, so these tests run in reverse mode only. +constexpr bool reverse_only = true; +auto product_eval = [](const auto& a, const auto& b) { return (a * b).eval(); }; +} // namespace + +TEST(mathMixMultiplyOperator, matrixMatrix) { + Eigen::MatrixXd a = test_matrix(3, 4, 0.1); + Eigen::MatrixXd b = test_matrix(4, 2, -0.2); + stan::test::expect_ad(product_eval, a, b); + Eigen::MatrixXd empty(0, 4); + stan::test::expect_ad(product_eval, empty, b); +} + +TEST(mathMixMultiplyOperator, matrixVector) { + Eigen::MatrixXd a = test_matrix(3, 4, 0.1); + Eigen::VectorXd v = test_matrix(4, 1, 0.3); + stan::test::expect_ad(product_eval, a, v); +} + +TEST(mathMixMultiplyOperator, rowVectorMatrix) { + Eigen::RowVectorXd rv = test_matrix(1, 3, 0.2); + Eigen::MatrixXd a = test_matrix(3, 4, 0.1); + stan::test::expect_ad(product_eval, rv, a); +} + +TEST(mathMixMultiplyOperator, innerAndOuterProduct) { + Eigen::RowVectorXd rv = test_matrix(1, 4, -0.1); + Eigen::VectorXd v = test_matrix(4, 1, 0.3); + stan::test::expect_ad(product_eval, rv, v); + stan::test::expect_ad(product_eval, v, rv); +} + +TEST(mathMixMultiplyOperator, lazyProduct) { + // lazy products are left to Eigen's coefficient-based evaluator + auto lazy_eval + = [](const auto& a, const auto& b) { return a.lazyProduct(b).eval(); }; + Eigen::MatrixXd a = test_matrix(2, 3, 0.1); + Eigen::MatrixXd b = test_matrix(3, 2, -0.2); + stan::test::expect_ad(lazy_eval, a, b); +} + +TEST(mathMixMultiplyOperator, expressionOperands) { + auto f_t = [](const auto& a, const auto& b) { + return (a.transpose() * b.block(0, 0, 3, 2)).eval(); + }; + Eigen::MatrixXd a = test_matrix(3, 4, 0.1); + Eigen::MatrixXd b = test_matrix(3, 3, -0.2); + stan::test::expect_ad(f_t, a, b); + auto f_chain = [](const auto& a, const auto& b) { + return (a * b * b.transpose()).eval(); + }; + stan::test::expect_ad(f_chain, a, test_matrix(4, 2, 0.4)); +} + +TEST(mathMixMultiplyOperator, assignment) { + auto f_add_sub = [](const auto& a, const auto& b) { + auto c = (a * b).eval(); + c += a * b; + c -= a * b.transpose(); + return c; + }; + auto f_noalias = [](const auto& a, const auto& b) { + Eigen::Matrix, -1, -1> c( + a.rows(), b.cols()); + c.noalias() = a * b; + c.noalias() += a * b; + return c; + }; + Eigen::MatrixXd a = test_matrix(3, 3, 0.1); + Eigen::MatrixXd b = test_matrix(3, 3, -0.2); + stan::test::expect_ad(f_add_sub, a, b); + stan::test::expect_ad(f_noalias, a, b); +} + +TEST(mathMixMultiplyOperator, triangularTimesVector) { + // Eigen's triangular kernels call the matrix-vector kernel directly. They + // do not support mixing autodiff and double scalars, so both operands are + // promoted to the common scalar. + auto f = [](const auto& a, const auto& v) { + using scalar = stan::return_type_t; + const Eigen::Matrix a_s = a; + const Eigen::Matrix v_s = v; + return (a_s.template triangularView() * v_s).eval(); + }; + Eigen::MatrixXd a = test_matrix(9, 9, 0.1); + Eigen::VectorXd v = test_matrix(9, 1, 0.3); + stan::test::expect_ad(f, a, v); +} diff --git a/test/unit/math/mix/fun/multiply_operator_scaled_test.cpp b/test/unit/math/mix/fun/multiply_operator_scaled_test.cpp new file mode 100644 index 00000000000..2a792024bc1 --- /dev/null +++ b/test/unit/math/mix/fun/multiply_operator_scaled_test.cpp @@ -0,0 +1,33 @@ +#include +#include + +// Values and derivatives of scaled matrix products `s * (A * B)`, +// `(A * B) * s` and `(s * A) * B`. Which Stan function computes each form is +// tested in test/unit/math/rev/core/eigen_product_hook_test.cpp. + +using stan::test::multiply_operator::test_matrix; + +TEST(mathMixMultiplyOperator, scaledProducts) { + Eigen::MatrixXd a = test_matrix(3, 2, 0.1); + Eigen::MatrixXd b = test_matrix(2, 2, -0.3); + double s = 0.7; + auto f_left = [](const auto& s, const auto& a, const auto& b) { + return (s * (a * b)).eval(); + }; + auto f_right = [](const auto& s, const auto& a, const auto& b) { + return ((a * b) * s).eval(); + }; + auto f_inner = [](const auto& s, const auto& a, const auto& b) { + return ((s * a) * (b * s)).eval(); + }; + auto f_sum = [](const auto& s, const auto& a, const auto& b) { + return (a.block(0, 0, 2, 2) + s * (a.block(0, 0, 2, 2) * b)).eval(); + }; + // Eigen cannot multiply fvar and double matrices until fvar products are + // routed to Stan's functions, so this test runs in reverse mode only. + constexpr bool reverse_only = true; + stan::test::expect_ad(f_left, s, a, b); + stan::test::expect_ad(f_right, s, a, b); + stan::test::expect_ad(f_inner, s, a, b); + stan::test::expect_ad(f_sum, s, a, b); +} diff --git a/test/unit/math/mix/fun/multiply_operator_util.hpp b/test/unit/math/mix/fun/multiply_operator_util.hpp new file mode 100644 index 00000000000..ec4a0883afc --- /dev/null +++ b/test/unit/math/mix/fun/multiply_operator_util.hpp @@ -0,0 +1,24 @@ +#ifndef TEST_UNIT_MATH_MIX_FUN_MULTIPLY_OPERATOR_UTIL_HPP +#define TEST_UNIT_MATH_MIX_FUN_MULTIPLY_OPERATOR_UTIL_HPP + +#include + +namespace stan { +namespace test { +namespace multiply_operator { + +inline Eigen::MatrixXd test_matrix(int rows, int cols, double offset) { + Eigen::MatrixXd x(rows, cols); + for (int j = 0; j < cols; ++j) { + for (int i = 0; i < rows; ++i) { + x(i, j) = 0.3 * (i + 1) - 0.2 * (j + 1) + offset; + } + } + return x; +} + +} // namespace multiply_operator +} // namespace test +} // namespace stan + +#endif diff --git a/test/unit/math/mix/fun/multiply_sub_test.cpp b/test/unit/math/mix/fun/multiply_sub_test.cpp new file mode 100644 index 00000000000..22bddc5de3d --- /dev/null +++ b/test/unit/math/mix/fun/multiply_sub_test.cpp @@ -0,0 +1,21 @@ +#include + +TEST(mathMixMatFun, multiplySub) { + auto f = [](const auto& a, const auto& b, const auto& c) { + return stan::math::multiply_sub(a, b, c); + }; + Eigen::MatrixXd a(3, 2); + a << 1, -2, 0.5, 3, -1.5, 2; + Eigen::MatrixXd b(2, 2); + b << 0.25, -1, 2, 0.5; + Eigen::MatrixXd c(3, 2); + c << 0.1, 0.2, -0.3, 0.4, 0.5, -0.6; + Eigen::VectorXd v(2); + v << -0.5, 1.5; + Eigen::VectorXd cv(3); + cv << 1, 2, 3; + // reverse mode only: multiply_sub has no fvar overload yet + constexpr bool reverse_only = true; + stan::test::expect_ad(f, a, b, c); + stan::test::expect_ad(f, a, v, cv); +} diff --git a/test/unit/math/mix/fun/scale_multiply_test.cpp b/test/unit/math/mix/fun/scale_multiply_test.cpp new file mode 100644 index 00000000000..0122006efaf --- /dev/null +++ b/test/unit/math/mix/fun/scale_multiply_test.cpp @@ -0,0 +1,23 @@ +#include + +TEST(mathMixMatFun, scaleMultiply) { + auto f = [](const auto& s, const auto& a, const auto& b) { + return stan::math::scale_multiply(s, a, b); + }; + double s = 0.7; + Eigen::MatrixXd a(3, 2); + a << 1, -2, 0.5, 3, -1.5, 2; + Eigen::MatrixXd b(2, 2); + b << 0.25, -1, 2, 0.5; + Eigen::VectorXd v(2); + v << -0.5, 1.5; + Eigen::RowVectorXd rv(3); + rv << 1, 0.5, -2; + // reverse mode only: scale_multiply has no fvar overload yet + constexpr bool reverse_only = true; + stan::test::expect_ad(f, s, a, b); + stan::test::expect_ad(f, s, a, v); + stan::test::expect_ad(f, s, rv, a); + Eigen::MatrixXd empty(0, 2); + stan::test::expect_ad(f, s, empty, b); +} diff --git a/test/unit/math/rev/core/eigen_cwise_product_evaluator_mre_test.cpp b/test/unit/math/rev/core/eigen_cwise_product_evaluator_mre_test.cpp new file mode 100644 index 00000000000..2eda1b4663a --- /dev/null +++ b/test/unit/math/rev/core/eigen_cwise_product_evaluator_mre_test.cpp @@ -0,0 +1,15 @@ +#include +#include +#include + +// Minimal instantiation of Eigen's binary_evaluator for a coefficient-wise +// product of two var arrays. `a * a` is a +// `CwiseBinaryOp, ...>`, and `.sum()` builds its +// evaluator directly. +TEST_F(AgradRev, eigen_cwise_product_evaluator_var_var) { + using stan::math::var; + Eigen::Array a(2, 2); + a << 1, 2, 3, 4; + const var var_var = (a * a).sum(); + EXPECT_DOUBLE_EQ(var_var.val(), 30.0); +} diff --git a/test/unit/math/rev/core/eigen_lazy_product_evaluator_mre_test.cpp b/test/unit/math/rev/core/eigen_lazy_product_evaluator_mre_test.cpp new file mode 100644 index 00000000000..3b5618a1c4c --- /dev/null +++ b/test/unit/math/rev/core/eigen_lazy_product_evaluator_mre_test.cpp @@ -0,0 +1,15 @@ +#include +#include +#include + +// Minimal instantiation of Eigen's product_evaluator for a lazy product of +// two var matrices. `A.lazyProduct(A)` is a `Product<..., LazyProduct>`, and +// `.sum()` builds its evaluator directly. Stan does not specialize lazy +// products, so this is Eigen's coefficient-based evaluator. +TEST_F(AgradRev, eigen_lazy_product_evaluator_var_var) { + using stan::math::var; + Eigen::Matrix a(2, 2); + a << 1, 2, 3, 4; + const var var_var = a.lazyProduct(a).sum(); + EXPECT_DOUBLE_EQ(var_var.val(), (a.val() * a.val()).sum()); +} diff --git a/test/unit/math/rev/core/eigen_product_assignment_mre_test.cpp b/test/unit/math/rev/core/eigen_product_assignment_mre_test.cpp new file mode 100644 index 00000000000..b0222585241 --- /dev/null +++ b/test/unit/math/rev/core/eigen_product_assignment_mre_test.cpp @@ -0,0 +1,16 @@ +#include +#include +#include + +// Minimal instantiation of Eigen's Assignment of a product of two var +// matrices. `c = a * a` assigns a `Product<...>` to a var matrix with +// `assign_op`, which Eigen handles in `Assignment` without a +// product_evaluator. +TEST_F(AgradRev, eigen_product_assignment_var_var) { + using stan::math::var; + Eigen::Matrix a(2, 2); + a << 1, 2, 3, 4; + Eigen::Matrix c(2, 2); + c = a * a; + EXPECT_DOUBLE_EQ(c.val().sum(), 54.0); +} diff --git a/test/unit/math/rev/core/eigen_product_evaluator_mre_test.cpp b/test/unit/math/rev/core/eigen_product_evaluator_mre_test.cpp new file mode 100644 index 00000000000..e9f3be49762 --- /dev/null +++ b/test/unit/math/rev/core/eigen_product_evaluator_mre_test.cpp @@ -0,0 +1,14 @@ +#include +#include +#include + +// Minimal instantiation of Eigen's product_evaluator for a product of two +// var matrices. `(A * A).sum()` builds the product's evaluator directly +// (unlike `C = A * A`, which goes through Eigen's Assignment). +TEST_F(AgradRev, eigen_product_evaluator_var_var) { + using stan::math::var; + Eigen::Matrix a(2, 2); + a << 1, 2, 3, 4; + const var var_var = (a * a).sum(); + EXPECT_DOUBLE_EQ(var_var.val(), (a.val() * a.val()).sum()); +} diff --git a/test/unit/math/rev/core/eigen_product_hook_test.cpp b/test/unit/math/rev/core/eigen_product_hook_test.cpp new file mode 100644 index 00000000000..06395b7e1be --- /dev/null +++ b/test/unit/math/rev/core/eigen_product_hook_test.cpp @@ -0,0 +1,455 @@ +#include +#include +#include +#include + +// Eigen's `operator*` on `var` matrices is evaluated by Stan's free functions +// (multiply, elt_multiply, scale_multiply, multiply_add, multiply_sub). Each +// test checks that an Eigen expression records the same nodes on the +// reverse-mode stack, and has the same value, as the direct call to the +// function it should use. + +namespace { + +using stan::math::var; +template +using mat = Eigen::Matrix; +template +using vec = Eigen::Matrix; +template +using row_vec = Eigen::Matrix; + +template +mat make_mat(Eigen::Index rows, Eigen::Index cols, double offset) { + mat x(rows, cols); + for (Eigen::Index j = 0; j < cols; ++j) { + for (Eigen::Index i = 0; i < rows; ++i) { + x(i, j) = 0.1 * (i + 1) - 0.07 * (j + 1) + offset; + } + } + return x; +} + +/** + * Expect `via_eigen()` to record the same chain and nochain varis as + * `direct()` and to return the same values. + */ +template +void expect_same_as_direct(const Direct& direct, const ViaEigen& via_eigen) { + const auto& stack = *stan::math::ChainableStack::instance_; + const std::size_t chain_start = stack.var_stack_.size(); + const std::size_t nochain_start = stack.var_nochain_stack_.size(); + const auto expected = direct(); + const std::size_t direct_chain = stack.var_stack_.size() - chain_start; + const std::size_t direct_nochain + = stack.var_nochain_stack_.size() - nochain_start; + const std::size_t chain_mid = stack.var_stack_.size(); + const std::size_t nochain_mid = stack.var_nochain_stack_.size(); + const auto actual = via_eigen(); + EXPECT_EQ(stack.var_stack_.size() - chain_mid, direct_chain); + EXPECT_EQ(stack.var_nochain_stack_.size() - nochain_mid, direct_nochain); + EXPECT_MATRIX_NEAR(stan::math::value_of(actual), + stan::math::value_of(expected), 1e-12); +} + +} // namespace + +TEST_F(AgradRev, eigen_product_hook_double_products_unchanged) { + const mat a = make_mat(3, 4, 0.2); + const mat b = make_mat(4, 2, -0.1); + const mat c = a * b; + const mat c_lazy = a.lazyProduct(b); + for (Eigen::Index j = 0; j < 2; ++j) { + for (Eigen::Index i = 0; i < 3; ++i) { + double expected = 0; + for (Eigen::Index k = 0; k < 4; ++k) { + expected += a(i, k) * b(k, j); + } + EXPECT_DOUBLE_EQ(c(i, j), expected); + EXPECT_DOUBLE_EQ(c_lazy(i, j), expected); + } + } + EXPECT_EQ(stan::math::ChainableStack::instance_->var_stack_.size(), 0); + EXPECT_EQ(stan::math::ChainableStack::instance_->var_nochain_stack_.size(), + 0); +} + +TEST_F(AgradRev, eigen_product_hook_matrix_products_use_multiply) { + mat a = make_mat(5, 4, 0.3); + mat b = make_mat(4, 3, -0.2); + const mat b_d = make_mat(4, 3, -0.2); + vec v = make_mat(4, 1, 0.1); + row_vec rv = make_mat(1, 5, 0.2); + expect_same_as_direct([&] { return stan::math::multiply(a, b); }, + [&] { return mat(a * b); }); + expect_same_as_direct([&] { return stan::math::multiply(a, b_d); }, + [&] { return mat(a * b_d); }); + expect_same_as_direct( + [&] { return stan::math::multiply(b_d.transpose(), a.transpose()); }, + [&] { return mat(b_d.transpose() * a.transpose()); }); + expect_same_as_direct([&] { return stan::math::multiply(a, v); }, + [&] { return vec(a * v); }); + expect_same_as_direct([&] { return stan::math::multiply(v, rv); }, + [&] { return mat(v * rv); }); + expect_same_as_direct( + [&] { + mat c(5, 3); + c.noalias() = stan::math::multiply(a, b); + return c; + }, + [&] { + mat c(5, 3); + c.noalias() = a * b; + return c; + }); +} + +TEST_F(AgradRev, eigen_product_hook_inner_product_uses_multiply) { + row_vec rv = make_mat(1, 4, 0.2); + vec v = make_mat(4, 1, 0.1); + const auto& stack = *stan::math::ChainableStack::instance_; + const std::size_t start + = stack.var_stack_.size() + stack.var_nochain_stack_.size(); + const var expected = stan::math::multiply(rv, v); + const std::size_t mid + = stack.var_stack_.size() + stack.var_nochain_stack_.size(); + const var actual = (rv * v)(0, 0); + EXPECT_EQ(stack.var_stack_.size() + stack.var_nochain_stack_.size() - mid, + mid - start); + EXPECT_NEAR(actual.val(), expected.val(), 1e-12); +} + +TEST_F(AgradRev, eigen_product_hook_cwise_products_use_elt_multiply) { + mat x = make_mat(4, 3, 0.3); + mat y = make_mat(4, 3, -0.4); + const mat y_d = make_mat(4, 3, -0.4); + expect_same_as_direct([&] { return stan::math::elt_multiply(x, y); }, + [&] { return mat(x.array() * y.array()); }); + expect_same_as_direct([&] { return stan::math::elt_multiply(x, y_d); }, + [&] { return mat(x.cwiseProduct(y_d)); }); +} + +TEST_F(AgradRev, eigen_product_hook_scalar_products_use_multiply) { + const mat a = make_mat(3, 2, 0.1); + mat a_v = make_mat(3, 2, 0.1); + var s = 1.5; + expect_same_as_direct([&] { return stan::math::multiply(s, a); }, + [&] { return mat(s * a); }); + expect_same_as_direct([&] { return stan::math::multiply(s, a); }, + [&] { return mat(a * s); }); + expect_same_as_direct([&] { return stan::math::multiply(2.0, a_v); }, + [&] { return mat(a_v * 2); }); + expect_same_as_direct([&] { return stan::math::multiply(s, a_v); }, + [&] { return mat((a_v.array() * s).matrix()); }); +} + +TEST_F(AgradRev, eigen_product_hook_scaled_products_use_scale_multiply) { + const mat a = make_mat(3, 4, 0.1); + mat b = make_mat(4, 2, -0.3); + var s = 0.7; + expect_same_as_direct([&] { return stan::math::scale_multiply(s, a, b); }, + [&] { return mat(s * (a * b)); }); + expect_same_as_direct([&] { return stan::math::scale_multiply(s, a, b); }, + [&] { return mat((a * b) * s); }); + expect_same_as_direct([&] { return stan::math::scale_multiply(s, a, b); }, + [&] { return mat((s * a) * b); }); + expect_same_as_direct([&] { return stan::math::scale_multiply(s, a, b); }, + [&] { return mat(a * (b * s)); }); +} + +TEST_F(AgradRev, eigen_product_hook_assignments_use_multiply_add) { + mat a = make_mat(3, 3, 0.3); + mat b = make_mat(3, 3, -0.1); + mat c = make_mat(3, 3, 0.05); + expect_same_as_direct([&] { return stan::math::multiply_add(a, b, c); }, + [&] { + mat d = c; + d.noalias() += a * b; + return d; + }); + expect_same_as_direct([&] { return stan::math::multiply_add(a, b, c); }, + [&] { return mat(c + a * b); }); + expect_same_as_direct([&] { return stan::math::multiply_sub(a, b, c); }, + [&] { + mat d = c; + d.noalias() -= a * b; + return d; + }); + // Without noalias() Eigen assumes the product may alias the destination, + // so it evaluates the product into a temporary first and then adds it + // coefficient by coefficient. + expect_same_as_direct( + [&] { + mat d = c; + d += mat(stan::math::multiply(a, b)); + return d; + }, + [&] { + mat d = c; + d += a * b; + return d; + }); +} + +TEST_F(AgradRev, eigen_product_hook_aliasing) { + mat a = make_mat(3, 3, 0.3); + mat b = make_mat(3, 3, -0.1); + mat e = a; + e = e * b; + EXPECT_MATRIX_NEAR(e.val(), a.val() * b.val(), 1e-12); + mat x = make_mat(3, 1, 0.2); + mat x0 = x; + x = x.cwiseProduct(x); + EXPECT_MATRIX_NEAR(x.val(), x0.val().cwiseProduct(x0.val()), 1e-12); +} + +TEST_F(AgradRev, eigen_product_hook_gradients) { + const Eigen::Index n = 3, k = 4, m = 2; + mat a = make_mat(n, k, 0.3); + mat b = make_mat(k, m, -0.2); + var s = 0.5; + mat c = s * (a * b); + stan::math::sum(c).grad(); + const mat ones = mat::Ones(n, m); + EXPECT_NEAR(s.adj(), (a.val() * b.val()).sum(), 1e-12); + EXPECT_MATRIX_NEAR(a.adj(), 0.5 * ones * b.val().transpose(), 1e-12); + EXPECT_MATRIX_NEAR(b.adj(), 0.5 * a.val().transpose() * ones, 1e-12); +} + +// A product nested inside a larger expression is still computed by Stan's +// function: each expression records the same tape as the same outer +// expression applied to the already computed product. + +TEST_F(AgradRev, eigen_product_hook_nested_views_of_products) { + mat a = make_mat(3, 4, 0.3); + mat b = make_mat(4, 3, -0.2); + mat c = make_mat(3, 3, 0.1); + const mat b_d = make_mat(4, 3, -0.2); + using stan::math::multiply; + expect_same_as_direct( + [&] { return mat(mat(multiply(a, b)).transpose()); }, + [&] { return mat((a * b).transpose()); }); + expect_same_as_direct( + [&] { return mat(mat(multiply(a, b)).block(0, 1, 2, 2)); }, + [&] { return mat((a * b).block(0, 1, 2, 2)); }); + expect_same_as_direct( + [&] { return mat(mat(multiply(a, b)).rowwise().sum()); }, + [&] { return mat((a * b).rowwise().sum()); }); + expect_same_as_direct( + [&] { + return mat( + (mat(multiply(a, b_d)) + c).block(1, 1, 2, 2).transpose()); + }, + [&] { return mat(((a * b_d) + c).block(1, 1, 2, 2).transpose()); }); + expect_same_as_direct( + [&] { + return mat(multiply(b.transpose(), a.transpose()).transpose()); + }, + [&] { return mat((b.transpose() * a.transpose()).transpose()); }); +} + +TEST_F(AgradRev, eigen_product_hook_nested_coefficient_wise) { + mat a = make_mat(3, 4, 0.3); + mat b = make_mat(4, 3, -0.2); + mat c = make_mat(3, 3, 0.1); + var s = 0.7; + using stan::math::multiply; + expect_same_as_direct( + [&] { return mat(mat(multiply(a, b)).array().exp().matrix()); }, + [&] { return mat((a * b).array().exp().matrix()); }); + expect_same_as_direct( + [&] { return mat(stan::math::elt_multiply(multiply(a, b), c)); }, + [&] { return mat((a * b).cwiseProduct(c)); }); + expect_same_as_direct( + [&] { + return mat( + mat(stan::math::scale_multiply(s, a, b)).transpose()); + }, + [&] { return mat((s * (a * b)).transpose()); }); + expect_same_as_direct( + [&] { return mat(mat(multiply(a, b)) + c - c); }, + [&] { return mat(a * b + c - c); }); +} + +TEST_F(AgradRev, eigen_product_hook_nested_products_of_products) { + mat a = make_mat(3, 4, 0.3); + mat b = make_mat(4, 3, -0.2); + mat c = make_mat(3, 3, 0.1); + vec v = make_mat(3, 1, 0.2); + using stan::math::multiply; + expect_same_as_direct([&] { return mat(multiply(multiply(a, b), c)); }, + [&] { return mat((a * b) * c); }); + expect_same_as_direct([&] { return mat(multiply(a, multiply(b, c))); }, + [&] { return mat(a * (b * c)); }); + expect_same_as_direct([&] { return vec(multiply(multiply(a, b), v)); }, + [&] { return vec(a * b * v); }); + expect_same_as_direct( + [&] { + return mat( + multiply(multiply(a, b), mat(multiply(a, b)).transpose())); + }, + [&] { return mat((a * b) * (a * b).transpose()); }); + expect_same_as_direct( + [&] { + return mat( + multiply(a.block(0, 0, 2, 3), b.topRows(3)).transpose()); + }, + [&] { + return mat((a.block(0, 0, 2, 3) * b.topRows(3)).transpose()); + }); +} + +TEST_F(AgradRev, eigen_product_hook_nested_in_stan_functions) { + mat a = make_mat(3, 4, 0.3); + mat b = make_mat(4, 3, -0.2); + mat c = make_mat(3, 3, 0.1); + using stan::math::multiply; + expect_same_as_direct( + [&] { return mat(stan::math::exp(mat(multiply(a, b)))); }, + [&] { return mat(stan::math::exp(a * b)); }); + expect_same_as_direct( + [&] { return mat(stan::math::add(mat(multiply(a, b)), c)); }, + [&] { return mat(stan::math::add(a * b, c)); }); + expect_same_as_direct( + [&] { + return mat::Constant(1, 1, + stan::math::sum(mat(multiply(a, b)))); + }, + [&] { return mat::Constant(1, 1, stan::math::sum(a * b)); }); +} + +TEST_F(AgradRev, eigen_product_hook_nested_destinations) { + mat a = make_mat(3, 4, 0.3); + mat b = make_mat(4, 3, -0.2); + mat c = make_mat(3, 3, 0.1); + using stan::math::multiply; + expect_same_as_direct( + [&] { + mat d = mat::Zero(5, 5); + d.block(1, 1, 3, 3) = multiply(a, b); + return d; + }, + [&] { + mat d = mat::Zero(5, 5); + d.block(1, 1, 3, 3) = a * b; + return d; + }); + expect_same_as_direct( + [&] { + mat d = mat::Zero(5, 5); + d.block(1, 1, 3, 3) = stan::math::multiply_add(a, b, c); + return d; + }, + [&] { + mat d = mat::Zero(5, 5); + d.block(1, 1, 3, 3) = c; + d.block(1, 1, 3, 3).noalias() += a * b; + return d; + }); + expect_same_as_direct( + [&] { + mat d = mat::Zero(3, 3); + d.transpose() = multiply(a, b); + return d; + }, + [&] { + mat d = mat::Zero(3, 3); + d.transpose() = a * b; + return d; + }); +} + +TEST_F(AgradRev, eigen_product_hook_nested_gradients) { + mat a = make_mat(3, 4, 0.3); + mat b = make_mat(4, 3, -0.2); + mat c = make_mat(3, 3, 0.1); + var s = 0.7; + stan::math::sum(((s * (a * b)).transpose().array() * c.array()).matrix() * c) + .grad(); + const mat a_v = a.val(); + const mat b_v = b.val(); + const mat c_v = c.val(); + // f = sum(((s A B)^T .* C) C) = sum_ij ((s A B)^T .* C)_ij r_j, r = C 1 + const Eigen::VectorXd r = c_v * Eigen::VectorXd::Ones(3); + // G = d f / d (s A B)^T = C .* (1 r^T); d f / d (s A B) = G^T + const mat g + = (c_v.array() * (Eigen::VectorXd::Ones(3) * r.transpose()).array()) + .matrix() + .transpose(); + EXPECT_NEAR(s.adj(), (g.array() * (a_v * b_v).array()).sum(), 1e-12); + EXPECT_MATRIX_NEAR(a.adj(), 0.7 * g * b_v.transpose(), 1e-12); + EXPECT_MATRIX_NEAR(b.adj(), 0.7 * a_v.transpose() * g, 1e-12); +} + +namespace { +/** + * Expect `op(x)` to propagate the same gradient to `x` when `x` is an + * `Eigen::Map` over a heap buffer of vars as when it is a plain matrix, even + * after the buffer is refilled with new vars before the reverse pass. Eigen's + * own algorithms (e.g. Householder reflections in `HouseholderQR`) map + * workspaces that they reuse and free like this. + */ +template +void expect_var_map_operand_kept(const Op& op) { + vec plain = make_mat(3, 1, 0.2); + stan::math::sum(op(plain)).grad(); + const Eigen::VectorXd expected = plain.adj(); + stan::math::set_zero_all_adjoints(); + + vec workspace = make_mat(3, 1, 0.2); + const vec forward_vars = workspace; + Eigen::Map> x(workspace.data(), workspace.size()); + const mat result = op(x); + for (Eigen::Index i = 0; i < workspace.size(); ++i) { + workspace(i) = var(10.0 + i); + } + stan::math::sum(result).grad(); + EXPECT_MATRIX_NEAR(forward_vars.adj(), expected, 1e-12); + EXPECT_MATRIX_NEAR(workspace.adj(), Eigen::VectorXd::Zero(3), 1e-12); +} +} // namespace + +TEST_F(AgradRev, eigen_product_hook_var_map_operands_are_copied) { + mat a = make_mat(3, 3, 0.3); + vec c = make_mat(3, 1, -0.1); + var s = 0.7; + // product, multiply + expect_var_map_operand_kept([&](const auto& x) { return mat(a * x); }); + // scaled product, scale_multiply + expect_var_map_operand_kept( + [&](const auto& x) { return mat((s * a) * x); }); + // noalias() +=, multiply_add + expect_var_map_operand_kept([&](const auto& x) { + mat d = c; + d.noalias() += a * x; + return d; + }); + // noalias() -=, multiply_sub + expect_var_map_operand_kept([&](const auto& x) { + mat d = c; + d.noalias() -= a * x; + return d; + }); + // scalar times matrix, multiply + expect_var_map_operand_kept([&](const auto& x) { return mat(s * x); }); + // coefficient-wise product, elt_multiply + expect_var_map_operand_kept( + [&](const auto& x) { return mat(x.cwiseProduct(c)); }); +} + +TEST_F(AgradRev, eigen_product_hook_var_map_destination_is_copied) { + mat a = make_mat(3, 3, 0.3); + vec b = make_mat(3, 1, -0.2); + // the destination of `noalias() +=` is an operand of multiply_add + vec buffer = make_mat(3, 1, 0.1); + const vec forward_vars = buffer; + Eigen::Map> dst(buffer.data(), buffer.size()); + dst.noalias() += a * b; + const vec result = buffer; + for (Eigen::Index i = 0; i < buffer.size(); ++i) { + buffer(i) = var(10.0 + i); + } + stan::math::sum(result).grad(); + EXPECT_MATRIX_NEAR(forward_vars.adj(), Eigen::VectorXd::Ones(3), 1e-12); + EXPECT_MATRIX_NEAR(buffer.adj(), Eigen::VectorXd::Zero(3), 1e-12); +} diff --git a/test/unit/math/rev/fun/eigen_hook_dispatch_test.cpp b/test/unit/math/rev/fun/eigen_hook_dispatch_test.cpp new file mode 100644 index 00000000000..f4e84129e96 --- /dev/null +++ b/test/unit/math/rev/fun/eigen_hook_dispatch_test.cpp @@ -0,0 +1,160 @@ +#include +#include +#include +#include +#include +#include + +// Checks which Eigen struct each product expression instantiates. Stan's +// specializations of Eigen's evaluators and assignments derive from +// `stan::math::internal::eigen_ad_hook`, so a type that derives from it is +// handled by Stan's functions. Each type is built as Eigen builds it: the +// expression types come from `decltype` of the expression, the assignment +// source is that type without `const` (Eigen deduces it from a `const Src&` +// argument), and the assignment functor is `Op`. + +namespace { + +using var = stan::math::var; +template +using mat = Eigen::Matrix; + +template +constexpr bool is_hooked + = std::is_base_of::value; + +template +constexpr bool evaluator_hooked = is_hooked>; + +template