From 9cbbe42aee11b7dd2de56371a33cdf73d34072d3 Mon Sep 17 00:00:00 2001 From: Steve Bronder Date: Wed, 30 Sep 2026 11:14:34 -0400 Subject: [PATCH 1/7] Evaluate Eigen products of autodiff scalars with Stan kernels Eigen `operator*` on matrices and arrays with `var`, `fvar`, `std::complex` or `std::complex>` scalars now dispatches to Stan kernels instead of Eigen's per-scalar kernels: - Matrix products (every product tag, lazyProduct, noalias, =, +=, -=) through a prim-level specialization of Eigen's dense product_evaluator that forwards autodiff products to eigen_matrix_product_impl and keeps Eigen's own evaluation for all other scalars, plus Assignment specializations for same-scalar destinations. - Coefficient-wise products and products with a scalar through binary_evaluator specializations for scalar_product_op. - Scalar factors are folded into the product they scale, so s * (A * B), (A * B) * s and (s * A) * B each record one callback. The reverse-mode kernel records one callback plus one nochain vari per result coefficient. The forward-mode kernel splits fvar products into products of values and tangents (one concatenated product for the tangent), so fvar matrix products record two callbacks instead of several varis per scalar operation; tape-free fvar element-wise products stay lazy. --- stan/math/fwd/core/eigen_product_hooks.hpp | 104 ++++ stan/math/fwd/core/eigen_product_kernel.hpp | 384 +++++++++++++ stan/math/fwd/core/fvar.hpp | 11 + stan/math/fwd/fun/Eigen_NumTraits.hpp | 2 + stan/math/prim/fun/Eigen.hpp | 2 + .../meta/eigen_product_evaluator_hooks.hpp | 515 ++++++++++++++++++ stan/math/rev/core/Eigen_NumTraits.hpp | 2 + stan/math/rev/core/eigen_product_hooks.hpp | 99 ++++ stan/math/rev/core/eigen_product_kernel.hpp | 425 +++++++++++++++ stan/math/rev/core/var.hpp | 11 + test/unit/math/ad_tape_util.hpp | 48 ++ .../math/fwd/core/eigen_product_hook_test.cpp | 82 +++ .../eigen_product_hook_kernels_tape_test.cpp | 83 +++ .../mix/core/eigen_product_hook_tape_test.cpp | 72 +++ .../fun/multiply_operator_complex_test.cpp | 49 ++ .../mix/fun/multiply_operator_cwise_test.cpp | 77 +++ .../mix/fun/multiply_operator_matmul_test.cpp | 128 +++++ .../mix/fun/multiply_operator_scaled_test.cpp | 48 ++ .../math/mix/fun/multiply_operator_util.hpp | 124 +++++ .../mix/prob/glm_forward_mode_tape_test.cpp | 10 +- .../math/rev/core/eigen_product_hook_test.cpp | 316 +++++++++++ 21 files changed, 2585 insertions(+), 7 deletions(-) create mode 100644 stan/math/fwd/core/eigen_product_hooks.hpp create mode 100644 stan/math/fwd/core/eigen_product_kernel.hpp create mode 100644 stan/math/prim/meta/eigen_product_evaluator_hooks.hpp create mode 100644 stan/math/rev/core/eigen_product_hooks.hpp create mode 100644 stan/math/rev/core/eigen_product_kernel.hpp create mode 100644 test/unit/math/ad_tape_util.hpp create mode 100644 test/unit/math/fwd/core/eigen_product_hook_test.cpp create mode 100644 test/unit/math/mix/core/eigen_product_hook_kernels_tape_test.cpp create mode 100644 test/unit/math/mix/core/eigen_product_hook_tape_test.cpp create mode 100644 test/unit/math/mix/fun/multiply_operator_complex_test.cpp create mode 100644 test/unit/math/mix/fun/multiply_operator_cwise_test.cpp create mode 100644 test/unit/math/mix/fun/multiply_operator_matmul_test.cpp create mode 100644 test/unit/math/mix/fun/multiply_operator_scaled_test.cpp create mode 100644 test/unit/math/mix/fun/multiply_operator_util.hpp create mode 100644 test/unit/math/rev/core/eigen_product_hook_test.cpp diff --git a/stan/math/fwd/core/eigen_product_hooks.hpp b/stan/math/fwd/core/eigen_product_hooks.hpp new file mode 100644 index 00000000000..50b0a394980 --- /dev/null +++ b/stan/math/fwd/core/eigen_product_hooks.hpp @@ -0,0 +1,104 @@ +#ifndef STAN_MATH_FWD_CORE_EIGEN_PRODUCT_HOOKS_HPP +#define STAN_MATH_FWD_CORE_EIGEN_PRODUCT_HOOKS_HPP + +#include +#include +#include +#include +#include +#include +#include + +/** + * @file + * Routes Eigen products whose operands have `fvar` or + * `std::complex>` scalars (mixed with arithmetic scalars) to the + * forward-mode kernels in stan/math/fwd/core/eigen_product_kernel.hpp, the + * same way stan/math/rev/core/eigen_product_hooks.hpp does for `var`. + */ + +namespace stan { +namespace math { +namespace internal { + +template +struct is_fwd_product_operand_scalar + : std::integral_constant::value> {}; +template +struct is_fwd_product_operand_scalar, T> + : std::integral_constant::value> {}; +template +struct is_fwd_product_operand_scalar, T> : std::true_type {}; +template +struct is_fwd_product_operand_scalar>, T> + : std::true_type {}; + +template +inline constexpr bool is_fwd_product_scalars_v + = !std::is_void::value && is_fwd_product_operand_scalar::value + && is_fwd_product_operand_scalar::value; + +/** + * The inner scalar `T` shared by the `fvar` operands of a product, or `void` + * if neither operand is an `fvar`. + */ +template +using fwd_product_inner_of_t + = std::conditional_t>::value, + fwd_product_inner_t, fwd_product_inner_t>; + +template +inline constexpr bool is_fwd_product_v = is_fwd_product_scalars_v< + typename Eigen::internal::traits::Scalar, + typename Eigen::internal::traits::Scalar, + fwd_product_inner_of_t::Scalar, + typename Eigen::internal::traits::Scalar>>; + +/** + * Forward-mode matrix product. + */ +template +struct eigen_matrix_product_impl>> { + using result_t = typename Eigen::Product::PlainObject; + static result_t run(const Lhs& lhs, const Rhs& rhs) { + return fwd_matrix_product(lhs, rhs); + } +}; + +/** + * Forward-mode coefficient-wise product, or product with a scalar. + */ +template +struct eigen_cwise_product_impl< + Eigen::CwiseBinaryOp, Lhs, Rhs>, + std::enable_if_t>> { + using xpr_t = Eigen::CwiseBinaryOp, + Lhs, Rhs>; + using result_t = typename xpr_t::PlainObject; + static auto run(const xpr_t& xpr) { return fwd_cwise_product(xpr); } +}; + +} // namespace internal +} // namespace math +} // namespace stan + +namespace Eigen { +namespace internal { + +#define STAN_FWD_PRODUCT_T_TPARAMS() typename T_, +#define STAN_FWD_PRODUCT_TS_TPARAMS() typename T_, typename S_, +STAN_EIGEN_AD_SCALAR_PATTERNS(STAN_EIGEN_AD_CWISE_PRODUCT_EVALUATOR, + STAN_FWD_PRODUCT_T_TPARAMS, + STAN_FWD_PRODUCT_TS_TPARAMS, stan::math::fvar) +STAN_EIGEN_AD_PRODUCT_ASSIGNMENT(STAN_FWD_PRODUCT_T_TPARAMS, + stan::math::fvar) +STAN_EIGEN_AD_PRODUCT_ASSIGNMENT(STAN_FWD_PRODUCT_T_TPARAMS, + std::complex>) +#undef STAN_FWD_PRODUCT_T_TPARAMS +#undef STAN_FWD_PRODUCT_TS_TPARAMS + +} // namespace internal +} // namespace Eigen + +#endif diff --git a/stan/math/fwd/core/eigen_product_kernel.hpp b/stan/math/fwd/core/eigen_product_kernel.hpp new file mode 100644 index 00000000000..9a65ae2189e --- /dev/null +++ b/stan/math/fwd/core/eigen_product_kernel.hpp @@ -0,0 +1,384 @@ +#ifndef STAN_MATH_FWD_CORE_EIGEN_PRODUCT_KERNEL_HPP +#define STAN_MATH_FWD_CORE_EIGEN_PRODUCT_KERNEL_HPP + +#include +#include +#include +#include +#include +#include +#include +#include +#include + +/** + * @file + * Forward-mode kernels for Eigen products whose operands have `fvar` or + * `std::complex>` scalars. A product of `fvar` matrices is split + * into products of their values and tangents, which are Eigen products of + * the inner scalar `T`. For `T = var` those go through the reverse-mode + * kernels, so an `fvar` matrix product puts two callbacks on the + * reverse-mode stack instead of several varis per scalar operation; for + * `T = double` they are ordinary double products. + * + * The kernels must never form a coefficient-wise or matrix product with + * `fvar` scalars themselves, since that would dispatch back into these + * kernels. Lazy results use `fwd_lazy_product_op` instead of Eigen's + * `scalar_product_op`. + */ + +namespace stan { +namespace math { +namespace internal { + +/** + * The inner scalar type `T` of `fvar` and `std::complex>`, and + * `void` for other scalars. + */ +template +struct fwd_product_inner { + using type = void; +}; +template +struct fwd_product_inner> { + using type = T; +}; +template +struct fwd_product_inner>> { + using type = T; +}; +template +using fwd_product_inner_t = typename fwd_product_inner::type; + +/** + * True if the innermost scalar of `fvar<......>` is a reverse-mode + * autodiff type, so products at that level record reverse-mode nodes. + */ +template +struct fwd_product_uses_tape : is_var {}; +template +struct fwd_product_uses_tape> : fwd_product_uses_tape {}; +template +struct fwd_product_uses_tape> : fwd_product_uses_tape {}; + +/** + * The value type used for an operand of a product with inner scalar `T`: + * `T` itself for `fvar`, `std::complex` for `std::complex>`, + * and `double` or `std::complex` for arithmetic operands. + */ +template +struct fwd_product_value { + using type = double; +}; +template +struct fwd_product_value, + std::enable_if_t::value>> { + using type = std::complex; +}; +template +struct fwd_product_value> { + using type = T; +}; +template +struct fwd_product_value>> { + using type = std::complex; +}; + +template +inline const T& fwd_product_val(const fvar& x) { + return x.val_; +} +template +inline const T& fwd_product_tangent(const fvar& x) { + return x.d_; +} +template +inline std::complex fwd_product_val(const std::complex>& z) { + return {z.real().val_, z.imag().val_}; +} +template +inline std::complex fwd_product_tangent(const std::complex>& z) { + return {z.real().d_, z.imag().d_}; +} + +template +inline const T& fwd_product_real(const T& x) { + return x; +} +template +inline T fwd_product_real(const std::complex& z) { + return z.real(); +} +template +inline double fwd_product_imag(const T&) { + return 0.0; +} +template +inline T fwd_product_imag(const std::complex& z) { + return z.imag(); +} + +/** + * Build a result coefficient from its value and tangent, which may have a + * lower scalar type (e.g. `double` for products of arithmetic operands + * scaled by an `fvar`). + */ +template +inline void fwd_product_assign(fvar& res, const V& val, const D& d) { + res = fvar(T(val), T(d)); +} +template +inline void fwd_product_assign(std::complex>& res, const V& val, + const D& d) { + res = std::complex>( + fvar(T(fwd_product_real(val)), T(fwd_product_real(d))), + fvar(T(fwd_product_imag(val)), T(fwd_product_imag(d)))); +} + +template +using fwd_operand_plain_t + = Eigen::Matrix; + +/** + * Operand of a forward-mode product kernel: its values and, for autodiff + * scalars, its tangents, each a plain matrix of the inner scalar type. + * + * @tparam X operand expression type + */ +template ::Scalar, + bool HasTangent = !std::is_void>::value> +struct fwd_product_operand { + static constexpr bool has_tangent = false; + using value_t = typename fwd_product_value::type; + fwd_operand_plain_t val_; + explicit fwd_product_operand(const X& x) : val_(x.template cast()) {} +}; + +template +struct fwd_product_operand { + static constexpr bool has_tangent = true; + using value_t = typename fwd_product_value::type; + fwd_operand_plain_t val_; + fwd_operand_plain_t d_; + explicit fwd_product_operand(const X& x) { + const fwd_operand_plain_t x_eval = x; + val_ = x_eval.unaryExpr( + [](const Scalar& s) -> value_t { return fwd_product_val(s); }); + d_ = x_eval.unaryExpr( + [](const Scalar& s) -> value_t { return fwd_product_tangent(s); }); + } +}; + +/** + * Scalar factors folded into a forward-mode product kernel, multiplied + * together as a scalar of the result type. + * + * @tparam ResScalar scalar type of the product result + */ +template +struct fwd_product_scales { + // empty until the first factor: constructing an `fvar`, even a zero + // or one, puts varis on the reverse-mode stack + std::optional value_; + bool any() const { return value_.has_value(); } + template + void push(const S& s) { + if (value_) { + *value_ *= s; + } else { + value_.emplace(s); + } + } +}; + +/** + * Write the value and tangent matrices of a kernel result into its plain + * result type. + */ +template +inline Result fwd_product_make_result(const Val& val, const Tangent& d) { + Result res(val.rows(), val.cols()); + for (Eigen::Index j = 0; j < val.cols(); ++j) { + for (Eigen::Index i = 0; i < val.rows(); ++i) { + fwd_product_assign(res.coeffRef(i, j), val.coeff(i, j), d.coeff(i, j)); + } + } + return res; +} + +/** + * The value and tangent of `sigma * M` from those of `M`, where `sigma` is + * the product of the folded scalar factors. + */ +template +inline Result fwd_product_scaled_result(const Scales& scales, const MVal& m_val, + const MTangent& m_d) { + if (!scales.any()) { + return fwd_product_make_result(m_val, m_d); + } + const auto s_val = fwd_product_val(*scales.value_); + const auto s_d = fwd_product_tangent(*scales.value_); + const auto val = (s_val * m_val).eval(); + const auto d = (s_d * m_val + s_val * m_d).eval(); + return fwd_product_make_result(val, d); +} + +/** + * `sigma * X * Y` for matrix operands. The tangent is `X.val * Y.d + X.d * + * Y.val`, computed as one product of the concatenations `[X.val, X.d]` and + * `[Y.d; Y.val]` when both operands have tangents. + */ +template +inline Result fwd_matmul(const Scales& scales, const XOp& x, const YOp& y) { + const auto m_val = (x.val_ * y.val_).eval(); + if constexpr (XOp::has_tangent && YOp::has_tangent) { + using x_value_t = typename XOp::value_t; + using y_value_t = typename YOp::value_t; + Eigen::Matrix x_cat( + x.val_.rows(), 2 * x.val_.cols()); + x_cat << x.val_, x.d_; + Eigen::Matrix y_cat( + 2 * y.val_.rows(), y.val_.cols()); + y_cat << y.d_, y.val_; + const auto m_d = (x_cat * y_cat).eval(); + return fwd_product_scaled_result(scales, m_val, m_d); + } else if constexpr (XOp::has_tangent) { + const auto m_d = (x.d_ * y.val_).eval(); + return fwd_product_scaled_result(scales, m_val, m_d); + } else if constexpr (YOp::has_tangent) { + const auto m_d = (x.val_ * y.d_).eval(); + return fwd_product_scaled_result(scales, m_val, m_d); + } else { + // arithmetic operands: the tangent comes from the scalar factors only + return fwd_product_scaled_result( + scales, m_val, + std::decay_t::Zero(m_val.rows(), m_val.cols())); + } +} + +/** + * `sigma * (X .* Y)` for operands of equal size. + */ +template +inline Result fwd_cwise(const Scales& scales, const XOp& x, const YOp& y) { + const auto m_val = (x.val_.array() * y.val_.array()).matrix().eval(); + if constexpr (XOp::has_tangent && YOp::has_tangent) { + const auto m_d + = (x.val_.array() * y.d_.array() + x.d_.array() * y.val_.array()) + .matrix() + .eval(); + return fwd_product_scaled_result(scales, m_val, m_d); + } else if constexpr (XOp::has_tangent) { + const auto m_d = (x.d_.array() * y.val_.array()).matrix().eval(); + return fwd_product_scaled_result(scales, m_val, m_d); + } else if constexpr (YOp::has_tangent) { + const auto m_d = (x.val_.array() * y.d_.array()).matrix().eval(); + return fwd_product_scaled_result(scales, m_val, m_d); + } else { + // arithmetic operands: the tangent comes from the scalar factors only + return fwd_product_scaled_result( + scales, m_val, + std::decay_t::Zero(m_val.rows(), m_val.cols())); + } +} + +/** + * `sigma * X` for a matrix operand. + */ +template +inline Result fwd_scale(const Scales& scales, const XOp& x) { + if constexpr (XOp::has_tangent) { + return fwd_product_scaled_result(scales, x.val_, x.d_); + } else { + return fwd_product_scaled_result( + scales, x.val_, + std::decay_t::Zero(x.val_.rows(), x.val_.cols()) + .eval()); + } +} + +/** + * Matrix product `lhs * rhs` of Eigen expressions with `fvar` scalars. + * + * @tparam Result plain Eigen type of the product + */ +template +inline Result fwd_matrix_product(const Lhs& lhs, const Rhs& rhs) { + fwd_product_scales scales; + const auto& lhs_core = peel_scales(scales, lhs); + const auto& rhs_core = peel_scales(scales, rhs); + using lhs_core_t = std::decay_t; + using rhs_core_t = std::decay_t; + return fwd_matmul(scales, fwd_product_operand(lhs_core), + fwd_product_operand(rhs_core)); +} + +/** + * Coefficient-wise product functor for lazy forward-mode products. It is + * not Eigen's `scalar_product_op`, so expressions using it are evaluated by + * Eigen's generic evaluator instead of Stan's product hooks. + */ +template +struct fwd_lazy_product_op { + using result_type = Result; + Result operator()(const L& a, const R& b) const { return a * b; } +}; + +/** + * Coefficient-wise product expression `xpr` (including a matrix times a + * scalar) with `fvar` scalars. + * + * When no reverse-mode tape is involved (e.g. `fvar`), splitting + * into value and tangent products would only add temporaries, so the + * product stays a lazy per-coefficient expression. Otherwise the value and + * tangent products are evaluated as products of the inner scalar, and a + * scaled matrix product `(A * B) * s` becomes one matrix-product kernel. + * + * @tparam Result plain Eigen type of the expression + */ +template +inline auto fwd_cwise_product(const Xpr& xpr) { + using res_scalar = typename Result::Scalar; + if constexpr (!fwd_product_uses_tape::value) { + using lhs_scalar = typename Eigen::internal::traits< + std::decay_t>::Scalar; + using rhs_scalar = typename Eigen::internal::traits< + std::decay_t>::Scalar; + return xpr.lhs().binaryExpr( + xpr.rhs(), fwd_lazy_product_op()); + } else { + fwd_product_scales scales; + const auto& core = peel_scales(scales, xpr); + using core_t = std::decay_t; + if constexpr (std::is_same::value) { + const auto& lhs_core = peel_scales(scales, xpr.lhs()); + const auto& rhs_core = peel_scales(scales, xpr.rhs()); + using lhs_core_t = std::decay_t; + using rhs_core_t = std::decay_t; + return fwd_cwise(scales, + fwd_product_operand(lhs_core), + fwd_product_operand(rhs_core)); + } else if constexpr (is_dense_product::value) { + const auto& lhs_core = peel_scales(scales, core.lhs()); + const auto& rhs_core = peel_scales(scales, core.rhs()); + using lhs_core_t = std::decay_t; + using rhs_core_t = std::decay_t; + return fwd_matmul(scales, + fwd_product_operand(lhs_core), + fwd_product_operand(rhs_core)); + } else { + return fwd_scale(scales, fwd_product_operand(core)); + } + } +} + +} // namespace internal +} // namespace math +} // namespace stan + +#endif diff --git a/stan/math/fwd/core/fvar.hpp b/stan/math/fwd/core/fvar.hpp index 187d6a5700b..cca3d751f87 100644 --- a/stan/math/fwd/core/fvar.hpp +++ b/stan/math/fwd/core/fvar.hpp @@ -1,6 +1,7 @@ #ifndef STAN_MATH_FWD_CORE_FVAR_HPP #define STAN_MATH_FWD_CORE_FVAR_HPP +#include #include #include #include @@ -256,6 +257,16 @@ struct fvar { } }; +namespace internal { +/** + * Eigen products of `fvar` scalars are evaluated by Stan (see + * stan/math/prim/meta/eigen_product_evaluator_hooks.hpp). Declared with + * `fvar` itself so every translation unit that uses `fvar` agrees. + */ +template +struct is_stan_ad_product_scalar> : std::true_type {}; +} // namespace internal + } // namespace math } // namespace stan #endif diff --git a/stan/math/fwd/fun/Eigen_NumTraits.hpp b/stan/math/fwd/fun/Eigen_NumTraits.hpp index 9e627e46d1e..9816fcc6d0d 100644 --- a/stan/math/fwd/fun/Eigen_NumTraits.hpp +++ b/stan/math/fwd/fun/Eigen_NumTraits.hpp @@ -253,4 +253,6 @@ struct ternary_evaluator< } // namespace internal } // namespace Eigen +#include + #endif diff --git a/stan/math/prim/fun/Eigen.hpp b/stan/math/prim/fun/Eigen.hpp index 9678be671ad..47f8d9867d0 100644 --- a/stan/math/prim/fun/Eigen.hpp +++ b/stan/math/prim/fun/Eigen.hpp @@ -87,4 +87,6 @@ } // namespace Eigen +#include + #endif diff --git a/stan/math/prim/meta/eigen_product_evaluator_hooks.hpp b/stan/math/prim/meta/eigen_product_evaluator_hooks.hpp new file mode 100644 index 00000000000..707bb40fcec --- /dev/null +++ b/stan/math/prim/meta/eigen_product_evaluator_hooks.hpp @@ -0,0 +1,515 @@ +#ifndef STAN_MATH_PRIM_META_EIGEN_PRODUCT_EVALUATOR_HOOKS_HPP +#define STAN_MATH_PRIM_META_EIGEN_PRODUCT_EVALUATOR_HOOKS_HPP + +/** + * @file + * Hooks that let Stan evaluate Eigen matrix products whose operands have + * autodiff scalars. + * + * Eigen dispatches `lhs * rhs` for dense operands through + * `Eigen::internal::product_evaluator`. Its scalar template parameters are + * defaulted from the operand types, so a partial specialization that fixes + * them is ambiguous with Eigen's generic specialization. Instead, the + * specializations below match every dense-by-dense product and choose, from + * the operand scalar types, between + * + * - `stan::math::internal::ad_product_evaluator`, which forwards to + * `eigen_matrix_product_impl::run(lhs, rhs)`, a customization + * point defined by the autodiff layers (`rev`, `fwd`); and + * - Eigen's own evaluation for all other scalars (e.g. `double`), which is + * reproduced exactly: `eigen_dense_product_evaluator` is a copy of the + * generic `product_evaluator` in Eigen/src/Core/ProductEvaluators.h, and + * lazy products derive from Eigen's coefficient-based evaluator. + * + * Because these specializations are selected for every dense product, they + * are included from stan/math/prim/fun/Eigen.hpp so that every translation + * unit sees the same definition. When upgrading Eigen, check the generic + * `product_evaluator` and the `LazyProduct` evaluator in + * Eigen/src/Core/ProductEvaluators.h against the copies here. + */ + +#include +#include +#include +#include + +namespace stan { +namespace math { +namespace internal { + +/** + * True for scalar types whose Eigen products Stan evaluates itself. The + * primary template is false; `rev` specializes it for `var` and `fwd` for + * `fvar`. Complex scalars inherit from their value type. + * + * @tparam Scalar scalar type of an Eigen expression + */ +template +struct is_stan_ad_product_scalar : std::false_type {}; + +template +struct is_stan_ad_product_scalar> + : is_stan_ad_product_scalar {}; + +/** + * True if either operand of an Eigen product has a scalar type for which + * Stan evaluates the product. + * + * @tparam Lhs left operand expression type + * @tparam Rhs right operand expression type + */ +template +inline constexpr bool is_stan_ad_product_v + = is_stan_ad_product_scalar< + typename Eigen::internal::traits::Scalar>::value + || is_stan_ad_product_scalar< + typename Eigen::internal::traits::Scalar>::value; + +/** + * Customization point for dense matrix products with autodiff operands. + * Specializations provide `static auto run(const Lhs&, const Rhs&)`, which + * returns an object with an Eigen evaluator holding the product, usually the + * product's plain object type. The primary template is intentionally left + * undefined so that a product with an autodiff scalar and no implementation + * is a compile error. + * + * @tparam Lhs left operand expression type + * @tparam Rhs right operand expression type + * @tparam Enable used by specializations for SFINAE + */ +template +struct eigen_matrix_product_impl; + +/** + * Customization point for coefficient-wise products (including the product + * of a matrix and a scalar) with autodiff operands. Specializations provide + * `static auto run(const Xpr&)` for the `CwiseBinaryOp` expression `Xpr`, + * returning either an evaluated plain object or a lazy expression; the + * implementation decides which. + * + * @tparam Xpr the `CwiseBinaryOp` of a `scalar_product_op` + * @tparam Enable used by specializations for SFINAE + */ +template +struct eigen_cwise_product_impl; + +/** + * Owns the object returned by a product customization point. It is a base + * class of `forwarding_evaluator` listed before the evaluator base, so the + * result is constructed before the evaluator that refers to it. + * + * @tparam Result type returned by the customization point + */ +template +struct forwarded_result { + Result result_; + template + explicit forwarded_result(F&& make) : result_(make()) {} +}; + +/** + * Eigen evaluator for an expression whose value is computed by a Stan + * customization point. The returned object is stored and read through its + * own Eigen evaluator. When the returned object is a plain matrix or array + * the evaluator sets `EvalBeforeNestingBit`, so an enclosing expression + * evaluates this one once instead of once per coefficient access. + * + * Eigen never copies evaluators (`evaluator_base` is noncopyable), so the + * inner evaluator's pointer into `result_` stays valid. + * + * @tparam Result type returned by the customization point + */ +template +struct forwarding_evaluator : private forwarded_result, + public Eigen::internal::evaluator { + using holder_t = forwarded_result; + using Base = Eigen::internal::evaluator; + static constexpr bool is_plain_result + = std::is_base_of, Result>::value; + enum { + Flags + = static_cast(Base::Flags) + | (is_plain_result ? static_cast(Eigen::EvalBeforeNestingBit) : 0) + }; + + template + explicit forwarding_evaluator(F&& make) + : holder_t(std::forward(make)), Base(this->holder_t::result_) {} +}; + +template +using matrix_product_result_t + = std::decay_t::run( + std::declval(), std::declval()))>; + +/** + * Evaluator of a dense product with an autodiff operand. It forwards to + * `eigen_matrix_product_impl`. + * + * @tparam Lhs left operand expression type + * @tparam Rhs right operand expression type + * @tparam Options Eigen product options (default, lazy, alias free) + */ +template +struct ad_product_evaluator + : forwarding_evaluator> { + using XprType = Eigen::Product; + using Base = forwarding_evaluator>; + static constexpr bool is_stan_ad_product_hook = true; + explicit ad_product_evaluator(const XprType& xpr) + : Base([&xpr]() { + return eigen_matrix_product_impl::run(xpr.lhs(), xpr.rhs()); + }) {} +}; + +/** + * Evaluator of a coefficient-wise product with an autodiff operand. It + * forwards to `eigen_cwise_product_impl`. + * + * @tparam Xpr the `CwiseBinaryOp` of a `scalar_product_op` + */ +template +struct ad_cwise_product_evaluator + : forwarding_evaluator::run(std::declval()))>> { + using XprType = Xpr; + using Base = forwarding_evaluator::run(std::declval()))>>; + static constexpr bool is_stan_ad_product_hook = true; + explicit ad_cwise_product_evaluator(const Xpr& xpr) + : Base([&xpr]() { return eigen_cwise_product_impl::run(xpr); }) {} +}; + +/** + * Detects the scalar-times-expression shapes Eigen builds for `s * X` and + * `X * s` (see `EIGEN_SCALAR_BINARYOP_EXPR_RETURN_TYPE` in Eigen's + * util/Macros.h), so a kernel can fold the scalar into one callback. + */ +template +struct scaled_expression : std::false_type {}; + +template +struct scaled_expression, + const Eigen::CwiseNullaryOp, Plain>, + const Inner>> : std::true_type { + template + static const S1& scale(const X& x) { + return x.lhs().functor().m_other; + } + template + static const Inner& inner(const X& x) { + return x.rhs(); + } +}; + +template +struct scaled_expression, const Inner, + const Eigen::CwiseNullaryOp, + Plain>>> : std::true_type { + template + static const S2& scale(const X& x) { + return x.rhs().functor().m_other; + } + template + static const Inner& inner(const X& x) { + return x.lhs(); + } +}; + +template +struct scaled_expression< + Eigen::CwiseBinaryOp, + const Eigen::CwiseNullaryOp< + Eigen::internal::scalar_constant_op, Plain1>, + const Eigen::CwiseNullaryOp< + Eigen::internal::scalar_constant_op, Plain2>>> + : std::true_type { + using Inner + = const Eigen::CwiseNullaryOp, + Plain2>; + template + static const S1& scale(const X& x) { + return x.lhs().functor().m_other; + } + template + static Inner& inner(const X& x) { + return x.rhs(); + } +}; + +/** + * Maximum number of nested scalar factors folded out of one operand. Deeper + * nests are left in the operand, which then goes through its own product + * kernel. + */ +constexpr int max_peel_depth = 4; + +/** + * Strip up to `max_peel_depth` scalar factors off `x` into `scales` and + * return the remaining expression. + */ +template +inline decltype(auto) peel_scales(Scales& scales, const X& x) { + if constexpr (scaled_expression::value && Depth < max_peel_depth) { + scales.push(scaled_expression::scale(x)); + return peel_scales(scales, scaled_expression::inner(x)); + } else { + return static_cast(x); + } +} + +/** + * True for an Eigen product of two dense operands. + */ +template +struct is_dense_product : std::false_type {}; +template +struct is_dense_product> + : std::integral_constant< + bool, + std::is_same::Shape, + Eigen::DenseShape>::value + && std::is_same< + typename Eigen::internal::evaluator_traits::Shape, + Eigen::DenseShape>::value> {}; + +/** + * True if the Eigen evaluator `Evaluator` is one of Stan's autodiff product + * evaluators, i.e. the expression it evaluates is computed by Stan. + * + * @tparam Evaluator an `Eigen::internal::evaluator` type + */ +template +struct is_stan_ad_product_hook : std::false_type {}; +template +struct is_stan_ad_product_hook< + Evaluator, std::void_t> + : std::integral_constant {}; + +/** + * True if Eigen evaluates the expression `Xpr` with one of Stan's autodiff + * product evaluators. + * + * @tparam Xpr an Eigen expression type + */ +template +inline constexpr bool is_stan_ad_product_hook_v + = is_stan_ad_product_hook>::value; + +/** + * Eigen's generic product evaluator for dense operands, reproduced from + * Eigen/src/Core/ProductEvaluators.h (`product_evaluator, ProductTag, LhsShape, RhsShape>`). It is used for products of + * scalars that Stan does not evaluate itself. + * + * @tparam Lhs left operand expression type + * @tparam Rhs right operand expression type + * @tparam Options Eigen product options + * @tparam ProductTag Eigen product implementation type + */ +template +struct eigen_dense_product_evaluator + : public Eigen::internal::evaluator< + typename Eigen::Product::PlainObject> { + using XprType = Eigen::Product; + using PlainObject = typename XprType::PlainObject; + using Base = Eigen::internal::evaluator; + enum { Flags = Base::Flags | Eigen::EvalBeforeNestingBit }; + + EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE explicit eigen_dense_product_evaluator( + const XprType& xpr) + : m_result(xpr.rows(), xpr.cols()) { + Eigen::internal::construct_at(this, m_result); + Eigen::internal::generic_product_impl::evalTo(m_result, + xpr.lhs(), + xpr.rhs()); + } + + protected: + PlainObject m_result; +}; + +/** + * `Assignment` of a product to a dense destination with the same autodiff + * scalar. Eigen's own specializations for this case call + * `generic_product_impl` directly and skip `product_evaluator`; for dense + * operands this forwards to `eigen_matrix_product_impl` instead. Products + * with non-dense operands (e.g. triangular views) run Eigen's body. + * + * @tparam Dst destination expression type + * @tparam Lhs left operand expression type + * @tparam Rhs right operand expression type + * @tparam Options Eigen product options + * @tparam Func assignment functor (`assign_op`, `add_assign_op`, + * `sub_assign_op`) + */ +template +struct ad_product_assignment { + using SrcXprType = Eigen::Product; + using Scalar = typename Dst::Scalar; + static constexpr bool dense_operands + = std::is_same::Shape, + Eigen::DenseShape>::value + && std::is_same::Shape, + Eigen::DenseShape>::value; + static void run(Dst& dst, const SrcXprType& src, const Func& func) { + if constexpr (dense_operands) { + const auto result + = eigen_matrix_product_impl::run(src.lhs(), src.rhs()); + Eigen::internal::call_assignment_no_alias(dst, result, func); + } else if constexpr (std::is_same>::value) { + if (dst.rows() != src.rows() || dst.cols() != src.cols()) { + dst.resize(src.rows(), src.cols()); + } + Eigen::internal::generic_product_impl::evalTo(dst, src.lhs(), + src.rhs()); + } else if constexpr (std::is_same>::value) { + Eigen::internal::generic_product_impl::addTo(dst, src.lhs(), + src.rhs()); + } else { + Eigen::internal::generic_product_impl::subTo(dst, src.lhs(), + src.rhs()); + } + } +}; + +/** + * Tag passed to Eigen's lazy (coefficient-based) product evaluator when it is + * reached from one of Stan's per-tag specializations. Eigen's lazy evaluator + * does not depend on its product tag, and no Eigen product has this tag, so + * it resolves to Eigen's own evaluator. + */ +constexpr int eigen_lazy_product_passthrough_tag = -1; + +template +using eigen_lazy_product_evaluator = Eigen::internal::product_evaluator< + Eigen::Product, + eigen_lazy_product_passthrough_tag, Eigen::DenseShape, Eigen::DenseShape>; + +} // namespace internal +} // namespace math +} // namespace stan + +namespace Eigen { +namespace internal { + +/** + * Dense-by-dense products: Stan's evaluator for autodiff operands, Eigen's + * generic evaluator otherwise. More specialized than Eigen's generic + * `product_evaluator` (fixed shapes) and less specialized than Eigen's lazy + * and `LazyCoeffBasedProductMode` evaluators, which are handled below. + */ +template +struct product_evaluator, ProductTag, DenseShape, + DenseShape> + : std::conditional_t< + stan::math::internal::is_stan_ad_product_v, + stan::math::internal::ad_product_evaluator, + stan::math::internal::eigen_dense_product_evaluator> { + using Base = std::conditional_t< + stan::math::internal::is_stan_ad_product_v, + stan::math::internal::ad_product_evaluator, + stan::math::internal::eigen_dense_product_evaluator>; + using Base::Base; +}; + +/** + * Lazy (coefficient-based) products, one specialization per product tag so + * each is more specialized than Eigen's `LazyProduct` evaluator. Eigen's + * `LazyCoeffBasedProductMode` evaluator for default products derives from + * the `CoeffBasedProductMode` one, so it is covered as well. + */ +#define STAN_EIGEN_LAZY_PRODUCT_EVALUATOR(TAG) \ + template \ + struct product_evaluator, TAG, DenseShape, \ + DenseShape> \ + : std::conditional_t< \ + stan::math::internal::is_stan_ad_product_v, \ + stan::math::internal::ad_product_evaluator, \ + stan::math::internal::eigen_lazy_product_evaluator> { \ + using Base = std::conditional_t< \ + stan::math::internal::is_stan_ad_product_v, \ + stan::math::internal::ad_product_evaluator, \ + stan::math::internal::eigen_lazy_product_evaluator>; \ + using Base::Base; \ + }; + +STAN_EIGEN_LAZY_PRODUCT_EVALUATOR(InnerProduct) +STAN_EIGEN_LAZY_PRODUCT_EVALUATOR(OuterProduct) +STAN_EIGEN_LAZY_PRODUCT_EVALUATOR(CoeffBasedProductMode) +STAN_EIGEN_LAZY_PRODUCT_EVALUATOR(LazyCoeffBasedProductMode) +STAN_EIGEN_LAZY_PRODUCT_EVALUATOR(GemvProduct) +STAN_EIGEN_LAZY_PRODUCT_EVALUATOR(GemmProduct) +#undef STAN_EIGEN_LAZY_PRODUCT_EVALUATOR + +} // namespace internal +} // namespace Eigen + +/** + * Apply `X(TPL, LhsScalar, RhsScalar)` to every scalar pattern of a product + * with an autodiff operand of type `AD` (and `std::complex`), where `S_` + * stands for any other scalar. `TPL` and `TPL_S` are function-like macros + * expanding to the extra template parameters (with a trailing comma) of the + * patterns without and with `S_`. The `(AD, AD)`, `(complex, + * complex)`, `(AD, complex)` and `(complex, AD)` patterns break + * the ties between the patterns with `S_`. + */ +#define STAN_EIGEN_AD_SCALAR_PATTERNS(X, TPL, TPL_S, AD) \ + X(TPL_S, AD, S_) \ + X(TPL_S, S_, AD) \ + X(TPL, AD, AD) \ + X(TPL_S, std::complex, S_) \ + X(TPL_S, S_, std::complex) \ + X(TPL, std::complex, std::complex) \ + X(TPL, AD, std::complex) \ + X(TPL, std::complex, AD) + +/** + * Specialize Eigen's `binary_evaluator` for coefficient-wise products with + * scalar pattern `(LS, RS)`, forwarding to `eigen_cwise_product_impl`. The + * scalar template parameters of `binary_evaluator` stay defaulted, exactly as + * in Eigen's generic specialization, so this one is more specialized only + * through the functor type. Use inside `namespace Eigen::internal`. + */ +#define STAN_EIGEN_AD_CWISE_PRODUCT_EVALUATOR(TPL, LS, RS) \ + template \ + struct binary_evaluator, Lhs, Rhs>, \ + IndexBased, IndexBased> \ + : stan::math::internal::ad_cwise_product_evaluator< \ + CwiseBinaryOp, Lhs, Rhs>> { \ + using Base = stan::math::internal::ad_cwise_product_evaluator< \ + CwiseBinaryOp, Lhs, Rhs>>; \ + using Base::Base; \ + }; + +/** + * Specialize Eigen's `Assignment` of a product to a dense destination with + * scalar `SCALAR` for `OP` (`assign_op`, `add_assign_op`, `sub_assign_op`). + * The last template argument repeats Eigen's own condition token for token, + * so this is more specialized only through the functor's scalar. Use inside + * `namespace Eigen::internal`. + */ +#define STAN_EIGEN_AD_PRODUCT_ASSIGNMENT_OP(TPL, SCALAR, OP) \ + template \ + struct Assignment, \ + OP, Dense2Dense, \ + std::enable_if_t<(Options == DefaultProduct \ + || Options == AliasFreeProduct)>> \ + : stan::math::internal::ad_product_assignment< \ + DstXprType, Lhs, Rhs, Options, OP> {}; + +#define STAN_EIGEN_AD_PRODUCT_ASSIGNMENT(TPL, SCALAR) \ + STAN_EIGEN_AD_PRODUCT_ASSIGNMENT_OP(TPL, SCALAR, assign_op) \ + STAN_EIGEN_AD_PRODUCT_ASSIGNMENT_OP(TPL, SCALAR, add_assign_op) \ + STAN_EIGEN_AD_PRODUCT_ASSIGNMENT_OP(TPL, SCALAR, sub_assign_op) + +#endif diff --git a/stan/math/rev/core/Eigen_NumTraits.hpp b/stan/math/rev/core/Eigen_NumTraits.hpp index ed5eebad4af..e3582600e6c 100644 --- a/stan/math/rev/core/Eigen_NumTraits.hpp +++ b/stan/math/rev/core/Eigen_NumTraits.hpp @@ -504,4 +504,6 @@ struct general_matrix_matrix_product + #endif diff --git a/stan/math/rev/core/eigen_product_hooks.hpp b/stan/math/rev/core/eigen_product_hooks.hpp new file mode 100644 index 00000000000..74a7127021d --- /dev/null +++ b/stan/math/rev/core/eigen_product_hooks.hpp @@ -0,0 +1,99 @@ +#ifndef STAN_MATH_REV_CORE_EIGEN_PRODUCT_HOOKS_HPP +#define STAN_MATH_REV_CORE_EIGEN_PRODUCT_HOOKS_HPP + +#include +#include +#include +#include +#include +#include +#include + +/** + * @file + * Routes Eigen products whose operands have `var` or `std::complex` + * scalars (mixed with arithmetic scalars) to the reverse-mode kernels in + * stan/math/rev/core/eigen_product_kernel.hpp. + * + * - matrix products `A * B` (every Eigen product tag, `lazyProduct`, + * `noalias()`, `=`, `+=` and `-=`) through the `product_evaluator` and + * `Assignment` hooks of + * stan/math/prim/meta/eigen_product_evaluator_hooks.hpp; + * - coefficient-wise products `a.array() * b.array()`, `cwiseProduct`, + * and products with a scalar `s * A`, `A * s` through `binary_evaluator`. + * + * Scalar factors are folded into the product they scale, so `s * (A * B)`, + * `(A * B) * s` and `(s * A) * B` each put one callback on the stack. + * + * The product is evaluated in full when its evaluator is created, even when + * only some coefficients are read afterwards (e.g. `(a * b).head(2)`). An + * unread coefficient's adjoint is zero, but the callback still multiplies it + * by the operand values, so a non-finite operand value can turn it into NaN. + * In particular, a product inside a branch of `.select()` is evaluated for + * every coefficient, unlike other expressions there. + */ + +namespace stan { +namespace math { +namespace internal { + +template +inline constexpr bool is_rev_product_v + = is_rev_product_operand_scalar< + typename Eigen::internal::traits::Scalar>::value&& + is_rev_product_operand_scalar< + typename Eigen::internal::traits::Scalar>::value + && (is_rev_product_scalar< + typename Eigen::internal::traits::Scalar>::value + || is_rev_product_scalar< + typename Eigen::internal::traits::Scalar>::value); + +/** + * Reverse-mode matrix product. + */ +template +struct eigen_matrix_product_impl>> { + using result_t = typename Eigen::Product::PlainObject; + static result_t run(const Lhs& lhs, const Rhs& rhs) { + return rev_matrix_product(lhs, rhs); + } +}; + +/** + * Reverse-mode coefficient-wise product, or product with a scalar. + */ +template +struct eigen_cwise_product_impl< + Eigen::CwiseBinaryOp, Lhs, Rhs>, + std::enable_if_t>> { + using xpr_t = Eigen::CwiseBinaryOp, + Lhs, Rhs>; + using result_t = typename xpr_t::PlainObject; + static result_t run(const xpr_t& xpr) { + return rev_cwise_product(xpr); + } +}; + +} // namespace internal +} // namespace math +} // namespace stan + +namespace Eigen { +namespace internal { + +#define STAN_REV_PRODUCT_NO_TPARAMS() +#define STAN_REV_PRODUCT_S_TPARAMS() typename S_, +STAN_EIGEN_AD_SCALAR_PATTERNS(STAN_EIGEN_AD_CWISE_PRODUCT_EVALUATOR, + STAN_REV_PRODUCT_NO_TPARAMS, + STAN_REV_PRODUCT_S_TPARAMS, stan::math::var) +STAN_EIGEN_AD_PRODUCT_ASSIGNMENT(STAN_REV_PRODUCT_NO_TPARAMS, stan::math::var) +STAN_EIGEN_AD_PRODUCT_ASSIGNMENT(STAN_REV_PRODUCT_NO_TPARAMS, + std::complex) +#undef STAN_REV_PRODUCT_NO_TPARAMS +#undef STAN_REV_PRODUCT_S_TPARAMS + +} // namespace internal +} // namespace Eigen + +#endif diff --git a/stan/math/rev/core/eigen_product_kernel.hpp b/stan/math/rev/core/eigen_product_kernel.hpp new file mode 100644 index 00000000000..204a267995f --- /dev/null +++ b/stan/math/rev/core/eigen_product_kernel.hpp @@ -0,0 +1,425 @@ +#ifndef STAN_MATH_REV_CORE_EIGEN_PRODUCT_KERNEL_HPP +#define STAN_MATH_REV_CORE_EIGEN_PRODUCT_KERNEL_HPP + +#include +#include +#include +#include +#include +#include +#include +#include +#include + +/** + * @file + * Reverse-mode kernels for Eigen products whose operands have `var` or + * `std::complex` scalars. Each product puts one callback on the + * reverse-mode stack plus one value-only (nochain) vari per result + * coefficient, instead of one vari per scalar operation. + * + * All arithmetic inside the kernels is on `double` or `std::complex` + * values. The kernels must never form an Eigen product of `var` scalars + * themselves, since that would dispatch back into these kernels. + */ + +namespace stan { +namespace math { +namespace internal { + +/** + * True for `var` and `std::complex`. + */ +template +struct is_rev_product_scalar : std::false_type {}; +template <> +struct is_rev_product_scalar : std::true_type {}; +template <> +struct is_rev_product_scalar> : std::true_type {}; + +/** + * True for scalars a reverse-mode product may mix with `var`. + */ +template +struct is_rev_product_operand_scalar + : bool_constant::value + || is_rev_product_scalar::value> {}; +template +struct is_rev_product_operand_scalar> + : bool_constant::value + || std::is_same::value> {}; + +/** + * The value type used for the arithmetic of a product whose result has + * scalar type `T`: `double`, or `std::complex` for complex results. + */ +template +struct rev_product_value { + using type = double; +}; +template +struct rev_product_value> { + using type = std::complex; +}; +template +using rev_product_value_t = typename rev_product_value::type; + +template * = nullptr> +inline double rev_product_value_of(const T& x) { + return x; +} +template * = nullptr> +inline std::complex rev_product_value_of(const std::complex& z) { + return {static_cast(z.real()), static_cast(z.imag())}; +} +inline double rev_product_value_of(const var& x) { return x.val(); } +inline std::complex rev_product_value_of(const std::complex& z) { + return {z.real().val(), z.imag().val()}; +} + +inline double rev_product_conj(double x) { return x; } +inline std::complex rev_product_conj(const std::complex& z) { + return std::conj(z); +} + +/** + * Add the adjoint `g` of a product result to operand `x`. For a real `x` and + * a complex `g`, only the real part contributes. + */ +inline void rev_product_add_adjoint(const var& x, double g) { + x.vi_->adj_ += g; +} +inline void rev_product_add_adjoint(const var& x, + const std::complex& g) { + x.vi_->adj_ += g.real(); +} +inline void rev_product_add_adjoint(const std::complex& z, + const std::complex& g) { + z.real().vi_->adj_ += g.real(); + z.imag().vi_->adj_ += g.imag(); +} + +/** + * A new value-only vari holding `x`. Results of a product kernel get their + * adjoints from the kernel's callback, so they need no `chain()`. + */ +inline var rev_product_result(double x) { return var(new vari(x, false)); } +inline std::complex rev_product_result(const std::complex& z) { + return {var(new vari(z.real(), false)), var(new vari(z.imag(), false))}; +} + +/** + * Operand of a reverse-mode product kernel: its values, and for autodiff + * scalars its vars, copied to the arena once. + * + * @tparam Scalar scalar type of the operand expression + */ +template ::value> +struct rev_product_operand { + static constexpr bool is_autodiff = false; + using value_t = rev_product_value_t; + arena_matrix> val_; + template + explicit rev_product_operand(const X& x) : val_(x.template cast()) {} + const auto& val() const { return val_; } + template + void add_adjoint(const G&) const {} +}; + +template +struct rev_product_operand { + static constexpr bool is_autodiff = true; + using value_t = rev_product_value_t; + arena_matrix> vars_; + arena_matrix> val_; + template + explicit rev_product_operand(const X& x) + : vars_(x), val_(vars_.unaryExpr([](const Scalar& s) { + return rev_product_value_of(s); + })) {} + const auto& val() const { return val_; } + template + void add_adjoint(const G& g) const { + const auto& g_ref = g.eval(); + for (Eigen::Index j = 0; j < vars_.cols(); ++j) { + for (Eigen::Index i = 0; i < vars_.rows(); ++i) { + rev_product_add_adjoint(vars_.coeff(i, j), g_ref.coeff(i, j)); + } + } + } +}; + +/** + * Scalar factors folded into a product kernel: the product of the + * arithmetic factors, and the autodiff factors. + * + * @tparam ValueT `double` or `std::complex` + */ +template +struct rev_product_scales { + // three peels (expression, left and right operand) of up to + // max_peel_depth factors each + static constexpr int max_autodiff = 12; + struct autodiff_scale { + vari* re_; + vari* im_; // nullptr for a real scale + ValueT value() const { + if constexpr (std::is_same::value) { + return re_->val_; + } else { + return {re_->val_, im_ == nullptr ? 0.0 : im_->val_}; + } + } + }; + ValueT constant_{1.0}; + autodiff_scale autodiff_[max_autodiff]; + int n_autodiff_{0}; + + template * = nullptr> + void push(const T& x) { + constant_ *= static_cast(x); + } + template * = nullptr> + void push(const std::complex& x) { + constant_ *= std::complex(x.real(), x.imag()); + } + void push(const var& x) { autodiff_[n_autodiff_++] = {x.vi_, nullptr}; } + void push(const std::complex& x) { + autodiff_[n_autodiff_++] = {x.real().vi_, x.imag().vi_}; + } + + /** Product of every factor. */ + ValueT value() const { + ValueT v = constant_; + for (int i = 0; i < n_autodiff_; ++i) { + v *= autodiff_[i].value(); + } + return v; + } + + /** + * Given `t = sum(G .* conj(M))` for the unscaled product `M` and result + * adjoints `G`, add each autodiff factor's adjoint. + */ + void add_adjoints(const ValueT& t) const { + for (int i = 0; i < n_autodiff_; ++i) { + ValueT others = constant_; + for (int k = 0; k < n_autodiff_; ++k) { + if (k != i) { + others *= autodiff_[k].value(); + } + } + const ValueT g = rev_product_conj(others) * t; + if constexpr (std::is_same::value) { + autodiff_[i].re_->adj_ += g; + } else { + autodiff_[i].re_->adj_ += g.real(); + if (autodiff_[i].im_ != nullptr) { + autodiff_[i].im_->adj_ += g.imag(); + } + } + } + } +}; + +template +using rev_result_value_t = rev_product_value_t; + +/** + * Build the plain result of a product kernel from its values, with one + * nochain vari per coefficient, and keep an arena copy for the callback. + */ +template +inline Result rev_product_make_result( + const Values& values, + arena_matrix>& res_arena) { + res_arena = arena_matrix< + Eigen::Matrix>( + values.rows(), values.cols()); + Result res(values.rows(), values.cols()); + for (Eigen::Index j = 0; j < values.cols(); ++j) { + for (Eigen::Index i = 0; i < values.rows(); ++i) { + res_arena.coeffRef(i, j) = rev_product_result(values.coeff(i, j)); + res.coeffRef(i, j) = res_arena.coeff(i, j); + } + } + return res; +} + +/** + * Adjoints of a kernel result as a `double` or `std::complex` + * matrix. + */ +template +inline auto rev_product_result_adjoints( + const arena_matrix< + Eigen::Matrix>& res) { + using value_t = rev_product_value_t; + Eigen::Matrix g(res.rows(), + res.cols()); + for (Eigen::Index j = 0; j < res.cols(); ++j) { + for (Eigen::Index i = 0; i < res.rows(); ++i) { + if constexpr (std::is_same::value) { + g.coeffRef(i, j) = res.coeff(i, j).adj(); + } else { + g.coeffRef(i, j) = value_t(res.coeff(i, j).real().adj(), + res.coeff(i, j).imag().adj()); + } + } + } + return g; +} + +template +using rev_operand_t + = rev_product_operand::Scalar>; + +/** + * `sigma * X * Y` for matrix operands `X` and `Y`, with one callback. + * + * @tparam Result plain Eigen type of the product + */ +template +inline Result rev_matmul(const Scales& scales, const XOp& x, const YOp& y) { + using value_t = rev_result_value_t; + using res_scalar = typename Result::Scalar; + const value_t sigma = scales.value(); + arena_matrix> + res_arena; + arena_matrix> unscaled; + Result res; + if (scales.n_autodiff_ > 0) { + unscaled = x.val() * y.val(); + res = rev_product_make_result((sigma * unscaled).eval(), res_arena); + } else { + res = rev_product_make_result((sigma * (x.val() * y.val())).eval(), + res_arena); + } + reverse_pass_callback([res_arena, x, y, scales, sigma, unscaled]() { + const auto g = rev_product_result_adjoints(res_arena); + const value_t sigma_conj = rev_product_conj(sigma); + if constexpr (XOp::is_autodiff) { + x.add_adjoint(sigma_conj * (g * y.val().adjoint())); + } + if constexpr (YOp::is_autodiff) { + y.add_adjoint(sigma_conj * (x.val().adjoint() * g)); + } + if (scales.n_autodiff_ > 0) { + scales.add_adjoints((g.array() * unscaled.array().conjugate()).sum()); + } + }); + return res; +} + +/** + * `sigma * (X .* Y)` for operands of equal size, with one callback. + */ +template +inline Result rev_cwise(const Scales& scales, const XOp& x, const YOp& y) { + using value_t = rev_result_value_t; + using res_scalar = typename Result::Scalar; + const value_t sigma = scales.value(); + arena_matrix> + res_arena; + Result res = rev_product_make_result( + (sigma * (x.val().array() * y.val().array())).matrix().eval(), res_arena); + reverse_pass_callback([res_arena, x, y, scales, sigma]() { + const auto g = rev_product_result_adjoints(res_arena); + const value_t sigma_conj = rev_product_conj(sigma); + if constexpr (XOp::is_autodiff) { + x.add_adjoint( + (sigma_conj * (g.array() * y.val().array().conjugate())).matrix()); + } + if constexpr (YOp::is_autodiff) { + y.add_adjoint( + (sigma_conj * (g.array() * x.val().array().conjugate())).matrix()); + } + if (scales.n_autodiff_ > 0) { + scales.add_adjoints( + (g.array() * (x.val().array() * y.val().array()).conjugate()).sum()); + } + }); + return res; +} + +/** + * `sigma * X` for a matrix operand `X`, with one callback. + */ +template +inline Result rev_scale(const Scales& scales, const XOp& x) { + using value_t = rev_result_value_t; + using res_scalar = typename Result::Scalar; + const value_t sigma = scales.value(); + arena_matrix> + res_arena; + Result res + = rev_product_make_result((sigma * x.val()).eval(), res_arena); + reverse_pass_callback([res_arena, x, scales, sigma]() { + const auto g = rev_product_result_adjoints(res_arena); + if constexpr (XOp::is_autodiff) { + x.add_adjoint(rev_product_conj(sigma) * g); + } + if (scales.n_autodiff_ > 0) { + scales.add_adjoints((g.array() * x.val().array().conjugate()).sum()); + } + }); + return res; +} + +/** + * Matrix product `lhs * rhs` of Eigen expressions. Scalar factors of either + * operand are folded into the product's callback. + * + * @tparam Result plain Eigen type of the product + */ +template +inline Result rev_matrix_product(const Lhs& lhs, const Rhs& rhs) { + rev_product_scales> scales; + const auto& lhs_core = peel_scales(scales, lhs); + const auto& rhs_core = peel_scales(scales, rhs); + using lhs_core_t = std::decay_t; + using rhs_core_t = std::decay_t; + return rev_matmul(scales, rev_operand_t(lhs_core), + rev_operand_t(rhs_core)); +} + +/** + * Coefficient-wise product expression `xpr` (including a matrix times a + * scalar). A scaled matrix product `(A * B) * s` becomes one matrix-product + * callback. + * + * @tparam Result plain Eigen type of the expression + */ +template +inline Result rev_cwise_product(const Xpr& xpr) { + using value_t = rev_result_value_t; + rev_product_scales scales; + const auto& core = peel_scales(scales, xpr); + using core_t = std::decay_t; + if constexpr (std::is_same::value) { + // a coefficient-wise product of two non-constant operands + const auto& lhs_core = peel_scales(scales, xpr.lhs()); + const auto& rhs_core = peel_scales(scales, xpr.rhs()); + using lhs_core_t = std::decay_t; + using rhs_core_t = std::decay_t; + return rev_cwise(scales, rev_operand_t(lhs_core), + rev_operand_t(rhs_core)); + } else if constexpr (is_dense_product::value) { + const auto& lhs_core = peel_scales(scales, core.lhs()); + const auto& rhs_core = peel_scales(scales, core.rhs()); + using lhs_core_t = std::decay_t; + using rhs_core_t = std::decay_t; + return rev_matmul(scales, rev_operand_t(lhs_core), + rev_operand_t(rhs_core)); + } else { + return rev_scale(scales, rev_operand_t(core)); + } +} + +} // namespace internal +} // namespace math +} // namespace stan + +#endif diff --git a/stan/math/rev/core/var.hpp b/stan/math/rev/core/var.hpp index 7186990ca98..8306fc8a2a0 100644 --- a/stan/math/rev/core/var.hpp +++ b/stan/math/rev/core/var.hpp @@ -8,6 +8,7 @@ #include #include #include +#include #include #include #include @@ -1186,6 +1187,16 @@ class var_value> { // For backwards compatibility the default value is double using var = var_value; +namespace internal { +/** + * Eigen products of `var` scalars are evaluated by Stan (see + * stan/math/prim/meta/eigen_product_evaluator_hooks.hpp). Declared with + * `var` itself so every translation unit that uses `var` agrees. + */ +template <> +struct is_stan_ad_product_scalar : std::true_type {}; +} // namespace internal + } // namespace math /** diff --git a/test/unit/math/ad_tape_util.hpp b/test/unit/math/ad_tape_util.hpp new file mode 100644 index 00000000000..d77fb6db502 --- /dev/null +++ b/test/unit/math/ad_tape_util.hpp @@ -0,0 +1,48 @@ +#ifndef TEST_UNIT_MATH_AD_TAPE_UTIL_HPP +#define TEST_UNIT_MATH_AD_TAPE_UTIL_HPP + +#include +#include + +namespace stan { +namespace test { + +/** + * Number of varis on the reverse-mode stack whose `chain()` runs in the + * reverse pass. + */ +inline std::size_t chain_stack_size() { + return stan::math::ChainableStack::instance_->var_stack_.size(); +} + +/** + * Number of varis on the reverse-mode stack that are skipped by the reverse + * pass (values without a chain method, e.g. results of a callback). + */ +inline std::size_t nochain_stack_size() { + return stan::math::ChainableStack::instance_->var_nochain_stack_.size(); +} + +/** + * Total number of varis on the reverse-mode stack. + */ +inline std::size_t tape_size() { + return chain_stack_size() + nochain_stack_size(); +} + +/** + * Records the stack sizes at construction and reports how many chain and + * nochain varis were added since. + */ +struct tape_delta { + std::size_t chain_start_{chain_stack_size()}; + std::size_t nochain_start_{nochain_stack_size()}; + std::size_t chain() const { return chain_stack_size() - chain_start_; } + std::size_t nochain() const { return nochain_stack_size() - nochain_start_; } + std::size_t total() const { return chain() + nochain(); } +}; + +} // namespace test +} // namespace stan + +#endif diff --git a/test/unit/math/fwd/core/eigen_product_hook_test.cpp b/test/unit/math/fwd/core/eigen_product_hook_test.cpp new file mode 100644 index 00000000000..a74d8b804a7 --- /dev/null +++ b/test/unit/math/fwd/core/eigen_product_hook_test.cpp @@ -0,0 +1,82 @@ +#include +#include +#include +#include + +// Eigen products of fvar matrices go through Stan's product hooks: +// matrix products are split into double products of values and tangents, +// coefficient-wise products stay lazy. + +namespace { +using stan::math::fvar; +using fvar_d = fvar; +using stan::math::internal::is_stan_ad_product_hook_v; + +template +Eigen::Matrix fvar_mat(int rows, int cols, double offset) { + Eigen::Matrix x(rows, cols); + for (int j = 0; j < cols; ++j) { + for (int i = 0; i < rows; ++i) { + x(i, j) = fvar_d(0.1 * (i + 1) - 0.2 * j + offset, 0.3 * j - 0.1 * i); + } + } + return x; +} + +template +void expect_fvar_product(const A& a, const B& b) { + Eigen::Matrix c = a * b; + ASSERT_EQ(c.rows(), a.rows()); + ASSERT_EQ(c.cols(), b.cols()); + for (int j = 0; j < c.cols(); ++j) { + for (int i = 0; i < c.rows(); ++i) { + fvar_d expected = 0; + for (int k = 0; k < a.cols(); ++k) { + expected += a(i, k) * b(k, j); + } + EXPECT_NEAR(c(i, j).val_, expected.val_, 1e-12); + EXPECT_NEAR(c(i, j).d_, expected.d_, 1e-12); + } + } +} +} // namespace + +TEST(mathFwdEigenProductHook, selected) { + using m = Eigen::Matrix; + using md = Eigen::MatrixXd; + EXPECT_TRUE((is_stan_ad_product_hook_v() * std::declval())>>)); + EXPECT_TRUE((is_stan_ad_product_hook_v() * std::declval())>>)); + EXPECT_TRUE((is_stan_ad_product_hook_v< + std::decay_t().array() + * std::declval().array())>>)); + EXPECT_FALSE((is_stan_ad_product_hook_v() * std::declval())>>)); +} + +TEST(mathFwdEigenProductHook, matmulValuesAndTangents) { + expect_fvar_product(fvar_mat(3, 4, 0.1), fvar_mat(4, 2, -0.2)); + const Eigen::MatrixXd d = Eigen::MatrixXd::Constant(4, 3, 0.5); + expect_fvar_product(fvar_mat(2, 4, 0.1), d); + expect_fvar_product(d.transpose(), fvar_mat(4, 1, 0.3)); + expect_fvar_product(fvar_mat<2, 2>(2, 2, 0.1), fvar_mat<2, 2>(2, 2, 0.4)); +} + +TEST(mathFwdEigenProductHook, cwiseIsLazy) { + auto a = fvar_mat(3, 2, 0.1); + auto b = fvar_mat(3, 2, -0.4); + auto xpr = a.array() * b.array(); + Eigen::internal::evaluator eval(xpr); + // lazy evaluation: the evaluator does not own a plain result + EXPECT_FALSE( + static_cast(Eigen::internal::evaluator::Flags + & Eigen::EvalBeforeNestingBit)); + for (int j = 0; j < 2; ++j) { + for (int i = 0; i < 3; ++i) { + const fvar_d expected = a(i, j) * b(i, j); + EXPECT_NEAR(eval.coeff(i, j).val_, expected.val_, 1e-12); + EXPECT_NEAR(eval.coeff(i, j).d_, expected.d_, 1e-12); + } + } +} 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..6e70ec935d2 --- /dev/null +++ b/test/unit/math/mix/core/eigen_product_hook_kernels_tape_test.cpp @@ -0,0 +1,83 @@ +#include +#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)); + } + } + kernel_chain_counts counts; + { + const stan::test::tape_delta delta; + Eigen::Matrix x + = a.template triangularView().solve(b); + counts.tri_solve = delta.chain(); + } + { + const stan::test::tape_delta delta; + Eigen::Matrix x = a.template triangularView() * b; + counts.tri_mul = delta.chain(); + } + { + const stan::test::tape_delta delta; + Eigen::Matrix x = a.llt().solve(b); + counts.llt = delta.chain(); + } + { + const stan::test::tape_delta delta; + Eigen::Matrix x = a.partialPivLu().solve(b); + counts.lu = delta.chain(); + } + 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/core/eigen_product_hook_tape_test.cpp b/test/unit/math/mix/core/eigen_product_hook_tape_test.cpp new file mode 100644 index 00000000000..f9e00b77902 --- /dev/null +++ b/test/unit/math/mix/core/eigen_product_hook_tape_test.cpp @@ -0,0 +1,72 @@ +#include +#include +#include +#include + +// Reverse-mode nodes recorded by Eigen products of fvar matrices, +// which Stan splits into products of their var values and tangents. + +namespace { +using stan::math::fvar; +using stan::math::var; +using fvar_v = fvar; + +Eigen::Matrix fvar_var_mat(int rows, int cols, double offset) { + Eigen::Matrix x(rows, cols); + for (int j = 0; j < cols; ++j) { + for (int i = 0; i < rows; ++i) { + x(i, j) = fvar_v(var(0.1 * (i + 1) - 0.2 * j + offset), var(1.0)); + } + } + return x; +} +} // namespace + +TEST(mathMixEigenProductHook, fvarVarProductsAreHooked) { + using stan::math::internal::is_stan_ad_product_hook_v; + using m = Eigen::Matrix; + using md = Eigen::MatrixXd; + EXPECT_TRUE((is_stan_ad_product_hook_v() * std::declval())>>)); + EXPECT_TRUE((is_stan_ad_product_hook_v() * std::declval())>>)); + EXPECT_TRUE((is_stan_ad_product_hook_v< + std::decay_t().array() + * std::declval().array())>>)); +} + +TEST(mathMixEigenProductHook, fvarVarMatmulTape) { + stan::math::nested_rev_autodiff nested; + const int n = 6, k = 5, m = 4; + auto a = fvar_var_mat(n, k, 0.1); + auto b = fvar_var_mat(k, m, -0.3); + const stan::test::tape_delta delta; + Eigen::Matrix c = a * b; + // one callback for the values and one for the concatenated tangent product + EXPECT_EQ(delta.chain(), 2); + EXPECT_EQ(delta.nochain(), 2 * n * m); +} + +TEST(mathMixEigenProductHook, fvarVarDoubleMatmulTape) { + stan::math::nested_rev_autodiff nested; + const int n = 6, k = 5; + const Eigen::MatrixXd x = Eigen::MatrixXd::Constant(n, k, 0.25); + auto beta = fvar_var_mat(k, 1, 0.2); + const stan::test::tape_delta delta; + Eigen::Matrix xb = x * beta; + EXPECT_EQ(delta.chain(), 2); + EXPECT_EQ(delta.nochain(), 2 * n); +} + +TEST(mathMixEigenProductHook, fvarVarCwiseTape) { + stan::math::nested_rev_autodiff nested; + const int n = 7; + auto a = fvar_var_mat(n, 1, 0.1); + auto b = fvar_var_mat(n, 1, -0.3); + const stan::test::tape_delta delta; + Eigen::Matrix c = (a.array() * b.array()).matrix(); + // values: one callback; tangent: two callbacks plus one vari per + // coefficient for their sum + EXPECT_EQ(delta.chain(), 3 + n); + EXPECT_EQ(delta.nochain(), 3 * n); +} diff --git a/test/unit/math/mix/fun/multiply_operator_complex_test.cpp b/test/unit/math/mix/fun/multiply_operator_complex_test.cpp new file mode 100644 index 00000000000..31d2ef05658 --- /dev/null +++ b/test/unit/math/mix/fun/multiply_operator_complex_test.cpp @@ -0,0 +1,49 @@ +#include +#include + +// Eigen products of complex autodiff matrices are evaluated by Stan's +// product hooks. + +using stan::test::multiply_operator::expect_hooked; + +namespace { +Eigen::MatrixXcd complex_matrix() { + Eigen::MatrixXcd x(2, 2); + x << std::complex(1, 0.5), 2, std::complex(3, -1), 4; + return x; +} +} // namespace + +TEST(mathMixMultiplyOperator, complexMatrixProduct) { + auto xpr = [](const auto& x, const auto& y) { return x * y; }; + auto f = [](const auto& x, const auto& y) { return (x * y).eval(); }; + Eigen::MatrixXd d_mat(2, 2); + d_mat << 1, 2, 3, 4; + Eigen::MatrixXcd cd_mat = complex_matrix(); + Eigen::VectorXcd cd_vec(2); + cd_vec << std::complex(0.5, 1), 2; + expect_hooked(xpr, cd_mat, cd_mat); + expect_hooked(xpr, cd_mat, cd_vec); + stan::test::expect_ad(f, d_mat, cd_mat); + stan::test::expect_ad(f, cd_mat, d_mat); + stan::test::expect_ad(f, cd_mat, cd_mat); + stan::test::expect_ad(f, cd_mat, cd_vec); +} + +TEST(mathMixMultiplyOperator, complexCwiseAndScalar) { + auto cwise = [](const auto& x, const auto& y) { return x.cwiseProduct(y); }; + auto scale = [](const auto& s, const auto& x) { return s * x; }; + auto f_cwise + = [](const auto& x, const auto& y) { return x.cwiseProduct(y).eval(); }; + auto f_scale = [](const auto& s, const auto& x) { return (s * x).eval(); }; + Eigen::MatrixXd d_mat(2, 2); + d_mat << 1, 2, 3, 4; + Eigen::MatrixXcd cd_mat = complex_matrix(); + std::complex cd_scalar(0.5, -0.25); + expect_hooked(cwise, cd_mat, cd_mat); + expect_hooked(scale, cd_scalar, cd_mat); + stan::test::expect_ad(f_cwise, cd_mat, cd_mat); + stan::test::expect_ad(f_cwise, d_mat, cd_mat); + stan::test::expect_ad(f_scale, cd_scalar, d_mat); + stan::test::expect_ad(f_scale, cd_scalar, cd_mat); +} 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..faece5b1c76 --- /dev/null +++ b/test/unit/math/mix/fun/multiply_operator_cwise_test.cpp @@ -0,0 +1,77 @@ +#include +#include + +// Coefficient-wise products and products with a scalar of autodiff +// matrices and arrays are evaluated by Stan's product hooks. + +using stan::test::multiply_operator::expect_hooked; +using stan::test::multiply_operator::test_matrix; + +TEST(mathMixMultiplyOperator, arrayTimesArray) { + auto xpr = [](const auto& a, const auto& b) { return a.array() * b.array(); }; + 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); + expect_hooked(xpr, a, b); + stan::test::expect_ad(f, a, b); +} + +TEST(mathMixMultiplyOperator, cwiseProduct) { + auto xpr = [](const auto& a, const auto& b) { return a.cwiseProduct(b); }; + 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); + expect_hooked(xpr, a, b); + 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; }; + auto right = [](const auto& s, const auto& a) { return a * s; }; + auto array = [](const auto& s, const auto& a) { return a.array() * s; }; + double s = 1.3; + Eigen::MatrixXd a = test_matrix(2, 3, 0.2); + expect_hooked(left, s, a); + expect_hooked(right, s, a); + expect_hooked(array, s, a); + auto eval = [](const auto& f) { + return [f](const auto& s, const auto& a) { + return Eigen::Matrix, -1, + -1>(f(s, a)); + }; + }; + stan::test::expect_ad(eval(left), s, a); + stan::test::expect_ad(eval(right), s, a); + stan::test::expect_ad(eval(array), s, a); +} + +TEST(mathMixMultiplyOperator, matrixTimesInt) { + Eigen::MatrixXd a = test_matrix(2, 3, 0.2); + // only the matrix can be autodiff here, so check the left promotion + using stan::math::fvar; + using stan::math::var; + EXPECT_TRUE( + (stan::math::internal::is_stan_ad_product_hook_v&>() * 3)>>)); + EXPECT_TRUE( + (stan::math::internal::is_stan_ad_product_hook_v, -1, -1>&>() * 3)>>)); + auto f = [](const auto& a) { return (a * 3).eval(); }; + stan::test::expect_ad(f, a); +} 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..f159db1c5e5 --- /dev/null +++ b/test/unit/math/mix/fun/multiply_operator_matmul_test.cpp @@ -0,0 +1,128 @@ +#include +#include + +// Eigen `operator*` matrix products of autodiff matrices are evaluated by +// Stan's product hooks. Each test checks that the hook is selected for every +// autodiff scalar and that values and derivatives are correct. + +using stan::test::multiply_operator::expect_assignment_hooked; +using stan::test::multiply_operator::expect_hooked; +using stan::test::multiply_operator::test_matrix; + +namespace { +auto product = [](const auto& a, const auto& b) { return a * b; }; +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); + expect_hooked(product, a, b); + 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); + expect_hooked(product, a, v); + 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); + expect_hooked(product, rv, a); + 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); + expect_hooked(product, rv, v); + expect_hooked(product, v, rv); + stan::test::expect_ad(product_eval, rv, v); + stan::test::expect_ad(product_eval, v, rv); +} + +TEST(mathMixMultiplyOperator, fixedSize) { + // 2x2 uses Eigen's lazy coefficient-based product, 4x4 the + // coefficient-based one + expect_hooked(product, Eigen::Matrix2d(test_matrix(2, 2, 0.1)), + Eigen::Matrix2d(test_matrix(2, 2, -0.3))); + expect_hooked(product, Eigen::Matrix4d(test_matrix(4, 4, 0.1)), + Eigen::Matrix4d(test_matrix(4, 4, -0.3))); + auto f_fixed = [](const auto& a, const auto& b) { + using a_scalar = typename std::decay_t::Scalar; + using b_scalar = typename std::decay_t::Scalar; + using ret_scalar = stan::return_type_t; + const Eigen::Matrix a2 = a; + const Eigen::Matrix b2 = b; + const Eigen::Matrix c = a2 * b2; + return Eigen::Matrix(c); + }; + stan::test::expect_ad(f_fixed, test_matrix(2, 2, 0.1), + test_matrix(2, 2, -0.3)); +} + +TEST(mathMixMultiplyOperator, lazyProduct) { + auto lazy = [](const auto& a, const auto& b) { return a.lazyProduct(b); }; + 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); + expect_hooked(lazy, a, b); + 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); + }; + auto f_t_eval + = [&f_t](const auto& a, const auto& b) { return f_t(a, b).eval(); }; + Eigen::MatrixXd a = test_matrix(3, 4, 0.1); + Eigen::MatrixXd b = test_matrix(3, 3, -0.2); + expect_hooked(f_t, a, b); + stan::test::expect_ad(f_t_eval, 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) { + expect_assignment_hooked(); + 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; + 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..5937bb9cd48 --- /dev/null +++ b/test/unit/math/mix/fun/multiply_operator_scaled_test.cpp @@ -0,0 +1,48 @@ +#include +#include + +// Scaled matrix products `s * (A * B)`, `(A * B) * s` and `(s * A) * B` are +// evaluated by Stan's product hooks with the scalar folded into the product. + +using stan::test::multiply_operator::expect_hooked; +using stan::test::multiply_operator::test_matrix; + +TEST(mathMixMultiplyOperator, scaledProductsHooked) { + const Eigen::MatrixXd b = test_matrix(2, 2, -0.3); + const Eigen::MatrixXd a = test_matrix(3, 2, 0.1); + double s = 0.7; + // autodiff scale and left operand, double right operand + expect_hooked([&b](const auto& s, const auto& a) { return s * (a * b); }, s, + a); + expect_hooked([&b](const auto& s, const auto& a) { return (a * b) * s; }, s, + a); + expect_hooked([&b](const auto& s, const auto& a) { return (s * a) * b; }, s, + a); + // autodiff scale and right operand, double left operand + expect_hooked([&a](const auto& s, const auto& b) { return s * (a * b); }, s, + b); + expect_hooked([&a](const auto& s, const auto& b) { return a * (b * s); }, s, + b); +} + +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(); + }; + 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..7b139907f6a --- /dev/null +++ b/test/unit/math/mix/fun/multiply_operator_util.hpp @@ -0,0 +1,124 @@ +#ifndef TEST_UNIT_MATH_MIX_FUN_MULTIPLY_OPERATOR_UTIL_HPP +#define TEST_UNIT_MATH_MIX_FUN_MULTIPLY_OPERATOR_UTIL_HPP + +#include +#include +#include +#include + +namespace stan { +namespace test { +namespace multiply_operator { + +using stan::math::fvar; +using stan::math::var; + +/** + * The autodiff scalars `expect_ad` instantiates, for which Eigen products + * must be evaluated by Stan's product hooks. + */ +using ad_scalars = std::tuple, fvar, fvar>, + fvar>>; + +template +struct promote { + using type = stan::math::promote_scalar_t; +}; +template +struct promote { + using type = T; +}; +template +struct promote> { + using type = std::complex; +}; +template +struct promote, R, C, O, MR, MC>> { + using type = Eigen::Matrix, R, C, O, MR, MC>; +}; +template +using promote_t = typename promote::type; + +/** + * Expect Eigen to evaluate `xpr(a, b)` with a Stan product hook when either + * or both arguments are promoted to autodiff scalar `T`. `xpr` returns the + * unevaluated product expression; it is only used in `decltype`. + */ +template +void expect_hooked_for(const F& xpr, const A&, const B&) { + using a_t = promote_t; + using b_t = promote_t; + using both_t = std::decay_t(), std::declval()))>; + using left_t = std::decay_t(), std::declval()))>; + using right_t = std::decay_t(), std::declval()))>; + EXPECT_TRUE(stan::math::internal::is_stan_ad_product_hook_v) + << "both " << typeid(T).name(); + EXPECT_TRUE(stan::math::internal::is_stan_ad_product_hook_v) + << "left " << typeid(T).name(); + EXPECT_TRUE(stan::math::internal::is_stan_ad_product_hook_v) + << "right " << typeid(T).name(); +} + +/** + * `expect_hooked_for` for every autodiff scalar, and expect the + * all-`double` expression to stay with Eigen. + */ +template +void expect_hooked(const F& xpr, const A& a, const B& b) { + std::apply( + [&](auto... t) { (expect_hooked_for(xpr, a, b), ...); }, + ad_scalars{}); + using double_t = std::decay_t; + EXPECT_FALSE(stan::math::internal::is_stan_ad_product_hook_v); +} + +/** + * Expect Eigen's assignment of a product to a destination with autodiff + * scalar `T` (`=`, `+=`, `-=`, `noalias()`) to use Stan's hook, which + * `Eigen::internal::Assignment` otherwise bypasses. + */ +template +void expect_assignment_hooked_for() { + using mat = Eigen::Matrix; + using product = Eigen::Product; + namespace ei = Eigen::internal; + EXPECT_TRUE((std::is_base_of>, + ei::Assignment, + ei::Dense2Dense>>::value)); + EXPECT_TRUE( + (std::is_base_of>, + ei::Assignment, + ei::Dense2Dense>>::value)); + EXPECT_TRUE( + (std::is_base_of>, + ei::Assignment, + ei::Dense2Dense>>::value)); +} + +inline void expect_assignment_hooked() { + std::apply( + [](auto... t) { (expect_assignment_hooked_for(), ...); }, + ad_scalars{}); +} + +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/prob/glm_forward_mode_tape_test.cpp b/test/unit/math/mix/prob/glm_forward_mode_tape_test.cpp index 30334e2ed36..ef35c573e4d 100644 --- a/test/unit/math/mix/prob/glm_forward_mode_tape_test.cpp +++ b/test/unit/math/mix/prob/glm_forward_mode_tape_test.cpp @@ -1,14 +1,10 @@ #include +#include #include #include namespace { -std::size_t tape_size() { - const auto& stack = *stan::math::ChainableStack::instance_; - return stack.var_stack_.size() + stack.var_nochain_stack_.size(); -} - // Nodes put on the reverse-mode tape by one forward-mode (fvar) sweep // of a GLM log density, the sweep hessian_times_vector runs for every // direction: the design-matrix products must not add one node per matrix @@ -27,10 +23,10 @@ std::size_t nodes_per_sweep(F&& f, int n, int k) { for (int i = 0; i < n; ++i) { alpha(i) = fvar(var(0.1), 1.0); } - const std::size_t before = tape_size(); + const stan::test::tape_delta delta; fvar lp = f(x, alpha, beta); EXPECT_TRUE(std::isfinite(lp.val_.val())); - return tape_size() - before; + return delta.total(); } } // namespace 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..b5fe5ad7f24 --- /dev/null +++ b/test/unit/math/rev/core/eigen_product_hook_test.cpp @@ -0,0 +1,316 @@ +#include +#include +#include +#include +#include +#include +#include + +namespace { + +using stan::math::var; +using stan::math::internal::is_stan_ad_product_hook_v; +using cvar = std::complex; +template +using mat = Eigen::Matrix; +template +using vec = Eigen::Matrix; +template +using row_vec = Eigen::Matrix; +template +using product_t = std::decay_t() + * std::declval())>; + +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; +} + +template +var sum_of(const Eigen::MatrixBase& x) { + var total = 0; + for (Eigen::Index j = 0; j < x.cols(); ++j) { + for (Eigen::Index i = 0; i < x.rows(); ++i) { + total += x.derived().coeff(i, j); + } + } + return total; +} + +} // namespace + +TEST_F(AgradRev, eigen_product_hook_selects_var_products) { + EXPECT_TRUE((is_stan_ad_product_hook_v, mat>>)); + EXPECT_TRUE((is_stan_ad_product_hook_v, mat>>)); + EXPECT_TRUE((is_stan_ad_product_hook_v, mat>>)); + EXPECT_TRUE( + (is_stan_ad_product_hook_v, mat>>)); + EXPECT_TRUE(( + is_stan_ad_product_hook_v, mat>>)); + EXPECT_TRUE((is_stan_ad_product_hook_v, vec>>)); + EXPECT_TRUE((is_stan_ad_product_hook_v, vec>>)); + EXPECT_TRUE((is_stan_ad_product_hook_v, row_vec>>)); + EXPECT_TRUE((is_stan_ad_product_hook_v, mat>>)); + EXPECT_TRUE((is_stan_ad_product_hook_v< + product_t>, mat>>)); + EXPECT_TRUE( + (is_stan_ad_product_hook_v< + std::decay_t&>().lazyProduct( + std::declval&>()))>>)); + EXPECT_FALSE( + (is_stan_ad_product_hook_v, mat>>)); + EXPECT_FALSE((is_stan_ad_product_hook_v< + product_t, mat>>)); +} + +TEST_F(AgradRev, eigen_product_hook_selects_var_cwise_products) { + using arr = Eigen::Array; + using arr_d = Eigen::Array; + EXPECT_TRUE((is_stan_ad_product_hook_v>)); + EXPECT_TRUE((is_stan_ad_product_hook_v>)); + EXPECT_TRUE((is_stan_ad_product_hook_v>)); + EXPECT_TRUE((is_stan_ad_product_hook_v>>)); + EXPECT_TRUE((is_stan_ad_product_hook_v, var>>)); + EXPECT_TRUE((is_stan_ad_product_hook_v, int>>)); + EXPECT_TRUE((is_stan_ad_product_hook_v>>)); + EXPECT_TRUE( + (is_stan_ad_product_hook_v< + std::decay_t&>().cwiseProduct( + std::declval&>()))>>)); + EXPECT_FALSE((is_stan_ad_product_hook_v>)); + EXPECT_FALSE((is_stan_ad_product_hook_v>>)); +} + +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::test::tape_size(), 0); +} + +TEST_F(AgradRev, eigen_product_hook_matmul_one_callback) { + const Eigen::Index n = 5, k = 4, m = 3; + mat a = make_mat(n, k, 0.3); + const mat b = make_mat(k, m, -0.2); + stan::test::tape_delta delta; + mat c = a * b; + EXPECT_EQ(delta.chain(), 1); + EXPECT_EQ(delta.nochain(), n * m); + const mat expected = a.val() * b; + EXPECT_MATRIX_NEAR(c.val(), expected, 1e-12); + sum_of(c).grad(); + // d sum(A * B) / dA = ones(n, m) * B^T + const mat grad_a = mat::Ones(n, m) * b.transpose(); + EXPECT_MATRIX_NEAR(a.adj(), grad_a, 1e-12); +} + +TEST_F(AgradRev, eigen_product_hook_matmul_var_var) { + 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); + stan::test::tape_delta delta; + mat c = a * b; + EXPECT_EQ(delta.chain(), 1); + EXPECT_EQ(delta.nochain(), n * m); + sum_of(c).grad(); + EXPECT_MATRIX_NEAR(a.adj(), mat::Ones(n, m) * b.val().transpose(), + 1e-12); + EXPECT_MATRIX_NEAR(b.adj(), a.val().transpose() * mat::Ones(n, m), + 1e-12); +} + +TEST_F(AgradRev, eigen_product_hook_matmul_shapes) { + vec v = make_mat(4, 1, 0.1); + row_vec rv = make_mat(1, 4, 0.2); + const mat a = make_mat(3, 4, 0.3); + { + stan::test::tape_delta delta; + vec av = a * v; + EXPECT_EQ(delta.chain(), 1); + EXPECT_EQ(delta.nochain(), 3); + EXPECT_MATRIX_NEAR(av.val(), a * v.val(), 1e-12); + } + { + stan::test::tape_delta delta; + var dot = (rv * v)(0, 0); + EXPECT_EQ(delta.chain(), 1); + EXPECT_EQ(delta.nochain(), 1); + EXPECT_NEAR(dot.val(), rv.val().dot(v.val()), 1e-12); + } + { + stan::test::tape_delta delta; + mat outer = v * rv; + EXPECT_EQ(delta.chain(), 1); + EXPECT_EQ(delta.nochain(), 16); + EXPECT_MATRIX_NEAR(outer.val(), v.val() * rv.val(), 1e-12); + } + { + mat f = make_mat(2, 2, 0.4); + mat g = make_mat(2, 2, 0.1); + stan::test::tape_delta delta; + mat fg = f * g; + mat fg_lazy = f.lazyProduct(g); + EXPECT_EQ(delta.chain(), 2); + EXPECT_EQ(delta.nochain(), 8); + EXPECT_MATRIX_NEAR(fg.val(), f.val() * g, 1e-12); + EXPECT_MATRIX_NEAR(fg_lazy.val(), f.val() * g, 1e-12); + } +} + +TEST_F(AgradRev, eigen_product_hook_assignment_forms) { + mat a = make_mat(3, 3, 0.3); + mat b = make_mat(3, 3, -0.1); + mat c = make_mat(3, 3, 0.05); + const mat c0 = c.val(); + { + stan::test::tape_delta delta; + mat d(3, 3); + d.noalias() = a * b; + EXPECT_EQ(delta.chain(), 1); + EXPECT_EQ(delta.nochain(), 9); + } + { + stan::test::tape_delta delta; + c += a * b; + // one callback for the product plus one vari per coefficient for += + EXPECT_EQ(delta.chain(), 1 + 9); + EXPECT_EQ(delta.nochain(), 9); + EXPECT_MATRIX_NEAR(c.val(), c0 + a.val() * b.val(), 1e-12); + } + { + // aliasing: the product is evaluated before the destination is written + mat e = a; + e = e * b; + EXPECT_MATRIX_NEAR(e.val(), a.val() * b.val(), 1e-12); + } +} + +TEST_F(AgradRev, eigen_product_hook_cwise_one_callback) { + const Eigen::Index n = 4; + vec x = make_mat(n, 1, 0.3); + vec y = make_mat(n, 1, -0.4); + stan::test::tape_delta delta; + vec z = (x.array() * y.array()).matrix(); + EXPECT_EQ(delta.chain(), 1); + EXPECT_EQ(delta.nochain(), n); + sum_of(z).grad(); + EXPECT_MATRIX_NEAR(x.adj(), y.val(), 1e-12); + EXPECT_MATRIX_NEAR(y.adj(), x.val(), 1e-12); +} + +TEST_F(AgradRev, eigen_product_hook_scalar_times_matrix) { + const mat a = make_mat(3, 2, 0.1); + var s = 1.5; + { + stan::test::tape_delta delta; + mat c = s * a; + EXPECT_EQ(delta.chain(), 1); + EXPECT_EQ(delta.nochain(), 6); + EXPECT_MATRIX_NEAR(c.val(), 1.5 * a, 1e-12); + sum_of(c).grad(); + EXPECT_NEAR(s.adj(), a.sum(), 1e-12); + } + stan::math::set_zero_all_adjoints(); + mat av = a; + { + stan::test::tape_delta delta; + mat c = av * 2; + EXPECT_EQ(delta.chain(), 1); + EXPECT_EQ(delta.nochain(), 6); + sum_of(c).grad(); + EXPECT_MATRIX_NEAR(av.adj(), mat::Constant(3, 2, 2.0), 1e-12); + } +} + +TEST_F(AgradRev, eigen_product_hook_scaled_products_fuse) { + const Eigen::Index n = 3, k = 4, m = 2; + const mat a = make_mat(n, k, 0.1); + mat b = make_mat(k, m, -0.3); + var s = 0.7; + const mat expected = 0.7 * (a * b.val()); + { + stan::test::tape_delta delta; + mat c = s * (a * b); + EXPECT_EQ(delta.chain(), 1); + EXPECT_EQ(delta.nochain(), n * m); + EXPECT_MATRIX_NEAR(c.val(), expected, 1e-12); + sum_of(c).grad(); + EXPECT_NEAR(s.adj(), (a * b.val()).sum(), 1e-12); + EXPECT_MATRIX_NEAR(b.adj(), 0.7 * a.transpose() * mat::Ones(n, m), + 1e-12); + } + { + stan::test::tape_delta delta; + mat c = (a * b) * s; + EXPECT_EQ(delta.chain(), 1); + EXPECT_EQ(delta.nochain(), n * m); + EXPECT_MATRIX_NEAR(c.val(), expected, 1e-12); + } + { + stan::test::tape_delta delta; + mat c = (s * a) * b; + EXPECT_EQ(delta.chain(), 1); + EXPECT_EQ(delta.nochain(), n * m); + EXPECT_MATRIX_NEAR(c.val(), expected, 1e-12); + } + { + stan::test::tape_delta delta; + mat c = a * (b * s); + EXPECT_EQ(delta.chain(), 1); + EXPECT_EQ(delta.nochain(), n * m); + EXPECT_MATRIX_NEAR(c.val(), expected, 1e-12); + } + { + mat d = make_mat(n, m, 0.2); + stan::test::tape_delta delta; + mat c = d + s * (a * b); + // one product callback plus one vari per coefficient for the sum + EXPECT_EQ(delta.chain(), 1 + n * m); + EXPECT_EQ(delta.nochain(), n * m); + EXPECT_MATRIX_NEAR(c.val(), d.val() + expected, 1e-12); + } +} + +TEST_F(AgradRev, eigen_product_hook_complex_matmul) { + mat a(2, 2); + a << cvar(1.0, 0.5), cvar(-0.3, 0.2), cvar(0.4, -1.0), cvar(2.0, 0.1); + mat> b(2, 1); + b << std::complex(0.5, -0.25), std::complex(1.5, 0.75); + stan::test::tape_delta delta; + mat c = a * b; + EXPECT_EQ(delta.chain(), 1); + EXPECT_EQ(delta.nochain(), 4); + mat> a_val(2, 2); + for (Eigen::Index j = 0; j < 2; ++j) { + for (Eigen::Index i = 0; i < 2; ++i) { + a_val(i, j) = {a(i, j).real().val(), a(i, j).imag().val()}; + } + } + const mat> expected = a_val * b; + for (Eigen::Index i = 0; i < 2; ++i) { + EXPECT_NEAR(c(i, 0).real().val(), expected(i, 0).real(), 1e-12); + EXPECT_NEAR(c(i, 0).imag().val(), expected(i, 0).imag(), 1e-12); + } + // d re(c_0) / d re(a_00) = re(b_0), d re(c_0) / d im(a_00) = -im(b_0) + c(0, 0).real().grad(); + EXPECT_NEAR(a(0, 0).real().adj(), b(0).real(), 1e-12); + EXPECT_NEAR(a(0, 0).imag().adj(), -b(0).imag(), 1e-12); + EXPECT_NEAR(a(1, 0).real().adj(), 0.0, 1e-12); +} From dcaea9ce0f95efb6881212e1d9a25297c9cffef0 Mon Sep 17 00:00:00 2001 From: Steve Bronder Date: Wed, 30 Sep 2026 11:17:49 -0400 Subject: [PATCH 2/7] Keep reverse-mode product callbacks small Scalar factors of a product kernel are stored on the arena at the first autodiff factor instead of in a fixed array inside every callback, result adjoints are read lazily, and var operand adjoints are updated with one vectorized expression. --- stan/math/rev/core/eigen_product_kernel.hpp | 63 ++++++++++++------- .../math/rev/core/eigen_product_hook_test.cpp | 24 +++++++ 2 files changed, 63 insertions(+), 24 deletions(-) diff --git a/stan/math/rev/core/eigen_product_kernel.hpp b/stan/math/rev/core/eigen_product_kernel.hpp index 204a267995f..a6b43a1cb87 100644 --- a/stan/math/rev/core/eigen_product_kernel.hpp +++ b/stan/math/rev/core/eigen_product_kernel.hpp @@ -141,10 +141,19 @@ struct rev_product_operand { const auto& val() const { return val_; } template void add_adjoint(const G& g) const { - const auto& g_ref = g.eval(); - for (Eigen::Index j = 0; j < vars_.cols(); ++j) { - for (Eigen::Index i = 0; i < vars_.rows(); ++i) { - rev_product_add_adjoint(vars_.coeff(i, j), g_ref.coeff(i, j)); + using g_scalar = typename Eigen::internal::traits::Scalar; + if constexpr (std::is_same::value + && std::is_same::value) { + // a shallow copy of the arena map, whose adjoint view is writable + arena_matrix> vars + = vars_; + vars.adj() += g; + } else { + const auto& g_ref = g.eval(); + for (Eigen::Index j = 0; j < vars_.cols(); ++j) { + for (Eigen::Index i = 0; i < vars_.rows(); ++i) { + rev_product_add_adjoint(vars_.coeff(i, j), g_ref.coeff(i, j)); + } } } } @@ -160,7 +169,7 @@ template struct rev_product_scales { // three peels (expression, left and right operand) of up to // max_peel_depth factors each - static constexpr int max_autodiff = 12; + static constexpr int max_autodiff = 3 * max_peel_depth; struct autodiff_scale { vari* re_; vari* im_; // nullptr for a real scale @@ -173,7 +182,9 @@ struct rev_product_scales { } }; ValueT constant_{1.0}; - autodiff_scale autodiff_[max_autodiff]; + // allocated on the arena at the first autodiff factor, so a kernel's + // callback captures a pointer instead of storage for every factor + autodiff_scale* autodiff_{nullptr}; int n_autodiff_{0}; template * = nullptr> @@ -184,9 +195,9 @@ struct rev_product_scales { void push(const std::complex& x) { constant_ *= std::complex(x.real(), x.imag()); } - void push(const var& x) { autodiff_[n_autodiff_++] = {x.vi_, nullptr}; } + void push(const var& x) { push_autodiff({x.vi_, nullptr}); } void push(const std::complex& x) { - autodiff_[n_autodiff_++] = {x.real().vi_, x.imag().vi_}; + push_autodiff({x.real().vi_, x.imag().vi_}); } /** Product of every factor. */ @@ -221,6 +232,16 @@ struct rev_product_scales { } } } + + private: + void push_autodiff(const autodiff_scale& s) { + if (autodiff_ == nullptr) { + autodiff_ + = ChainableStack::instance_->memalloc_.alloc_array( + max_autodiff); + } + autodiff_[n_autodiff_++] = s; + } }; template @@ -249,27 +270,20 @@ inline Result rev_product_make_result( } /** - * Adjoints of a kernel result as a `double` or `std::complex` - * matrix. + * Adjoints of a kernel result as a lazy `double` or `std::complex` + * expression. */ template inline auto rev_product_result_adjoints( const arena_matrix< Eigen::Matrix>& res) { - using value_t = rev_product_value_t; - Eigen::Matrix g(res.rows(), - res.cols()); - for (Eigen::Index j = 0; j < res.cols(); ++j) { - for (Eigen::Index i = 0; i < res.rows(); ++i) { - if constexpr (std::is_same::value) { - g.coeffRef(i, j) = res.coeff(i, j).adj(); - } else { - g.coeffRef(i, j) = value_t(res.coeff(i, j).real().adj(), - res.coeff(i, j).imag().adj()); - } - } + if constexpr (std::is_same::value) { + return res.adj(); + } else { + return res.unaryExpr([](const ResScalar& z) { + return std::complex(z.real().adj(), z.imag().adj()); + }); } - return g; } template @@ -298,7 +312,8 @@ inline Result rev_matmul(const Scales& scales, const XOp& x, const YOp& y) { res_arena); } reverse_pass_callback([res_arena, x, y, scales, sigma, unscaled]() { - const auto g = rev_product_result_adjoints(res_arena); + // evaluated once: both operand gradients are products with it + const auto g = rev_product_result_adjoints(res_arena).eval(); const value_t sigma_conj = rev_product_conj(sigma); if constexpr (XOp::is_autodiff) { x.add_adjoint(sigma_conj * (g * y.val().adjoint())); diff --git a/test/unit/math/rev/core/eigen_product_hook_test.cpp b/test/unit/math/rev/core/eigen_product_hook_test.cpp index b5fe5ad7f24..f4f0fa4176c 100644 --- a/test/unit/math/rev/core/eigen_product_hook_test.cpp +++ b/test/unit/math/rev/core/eigen_product_hook_test.cpp @@ -288,6 +288,30 @@ TEST_F(AgradRev, eigen_product_hook_scaled_products_fuse) { } } +TEST_F(AgradRev, eigen_product_hook_scales_are_small) { + // every product callback captures its scalar factors by value, so they + // must not embed storage for the factors themselves + EXPECT_LE(sizeof(stan::math::internal::rev_product_scales), + 4 * sizeof(void*)); + EXPECT_LE( + sizeof(stan::math::internal::rev_product_scales>), + 5 * sizeof(void*)); +} + +TEST_F(AgradRev, eigen_product_hook_many_autodiff_scales) { + const mat a = make_mat(2, 3, 0.1); + const mat b = make_mat(3, 2, -0.2); + var s1 = 0.5, s2 = 2.0, s3 = -1.5; + mat c = s1 * (s2 * a) * (b * s3); + const double scale = 0.5 * 2.0 * -1.5; + EXPECT_MATRIX_NEAR(c.val(), scale * (a * b), 1e-12); + sum_of(c).grad(); + const double total = (a * b).sum(); + EXPECT_NEAR(s1.adj(), 2.0 * -1.5 * total, 1e-12); + EXPECT_NEAR(s2.adj(), 0.5 * -1.5 * total, 1e-12); + EXPECT_NEAR(s3.adj(), 0.5 * 2.0 * total, 1e-12); +} + TEST_F(AgradRev, eigen_product_hook_complex_matmul) { mat a(2, 2); a << cvar(1.0, 0.5), cvar(-0.3, 0.2), cvar(0.4, -1.0), cvar(2.0, 0.1); From 7e84424199ba0d6d8b78edbca1bafaedd2e782aa Mon Sep 17 00:00:00 2001 From: Steve Bronder Date: Wed, 30 Sep 2026 13:14:18 -0400 Subject: [PATCH 3/7] Route Eigen var products to Stan's functions from Eigen's own structs Replace the evaluator/customization-point layers and the rev and fwd product kernels with specializations of Eigen's structs that call Stan's free functions directly: - product_evaluator (dense and per-tag lazy) calls multiply, or scale_multiply when an operand is Eigen's `s * A` / `A * s`. - binary_evaluator for scalar_product_op of var with var, double or int calls elt_multiply, multiply(s, X) or scale_multiply. - Assignment of a product to a var destination calls multiply, multiply_add (`noalias() +=`, `C + A * B`) or subtract (`noalias() -=`). Add scale_multiply(s, A, B) and multiply_add(A, B, C) (prim and rev), each recording one callback. This pass handles real var only: products with fvar or complex scalars use Eigen's evaluation, and the tests that need them are reverse-only or removed until those are routed as well. --- stan/math/fwd/core/eigen_product_hooks.hpp | 104 --- stan/math/fwd/core/eigen_product_kernel.hpp | 384 ---------- stan/math/fwd/core/fvar.hpp | 11 - stan/math/fwd/fun/Eigen_NumTraits.hpp | 1 - stan/math/prim/fun.hpp | 2 + stan/math/prim/fun/multiply_add.hpp | 41 + stan/math/prim/fun/scale_multiply.hpp | 36 + .../meta/eigen_product_evaluator_hooks.hpp | 713 +++++++----------- stan/math/rev/core/Eigen_NumTraits.hpp | 24 +- stan/math/rev/core/eigen_product_hooks.hpp | 99 --- stan/math/rev/core/eigen_product_kernel.hpp | 440 ----------- stan/math/rev/core/var.hpp | 11 - stan/math/rev/fun.hpp | 2 + stan/math/rev/fun/multiply_add.hpp | 69 ++ stan/math/rev/fun/scale_multiply.hpp | 66 ++ .../math/fwd/core/eigen_product_hook_test.cpp | 82 -- .../mix/core/eigen_product_hook_tape_test.cpp | 72 -- test/unit/math/mix/fun/multiply_add_test.cpp | 21 + .../fun/multiply_operator_complex_test.cpp | 49 -- .../mix/fun/multiply_operator_cwise_test.cpp | 45 +- .../mix/fun/multiply_operator_matmul_test.cpp | 59 +- .../mix/fun/multiply_operator_scaled_test.cpp | 35 +- .../math/mix/fun/multiply_operator_util.hpp | 100 --- .../unit/math/mix/fun/scale_multiply_test.cpp | 23 + .../math/rev/core/eigen_product_hook_test.cpp | 421 ++++------- test/unit/math/rev/fun/multiply_add_test.cpp | 48 ++ .../unit/math/rev/fun/scale_multiply_test.cpp | 47 ++ 27 files changed, 861 insertions(+), 2144 deletions(-) delete mode 100644 stan/math/fwd/core/eigen_product_hooks.hpp delete mode 100644 stan/math/fwd/core/eigen_product_kernel.hpp create mode 100644 stan/math/prim/fun/multiply_add.hpp create mode 100644 stan/math/prim/fun/scale_multiply.hpp delete mode 100644 stan/math/rev/core/eigen_product_hooks.hpp delete mode 100644 stan/math/rev/core/eigen_product_kernel.hpp create mode 100644 stan/math/rev/fun/multiply_add.hpp create mode 100644 stan/math/rev/fun/scale_multiply.hpp delete mode 100644 test/unit/math/fwd/core/eigen_product_hook_test.cpp delete mode 100644 test/unit/math/mix/core/eigen_product_hook_tape_test.cpp create mode 100644 test/unit/math/mix/fun/multiply_add_test.cpp delete mode 100644 test/unit/math/mix/fun/multiply_operator_complex_test.cpp create mode 100644 test/unit/math/mix/fun/scale_multiply_test.cpp create mode 100644 test/unit/math/rev/fun/multiply_add_test.cpp create mode 100644 test/unit/math/rev/fun/scale_multiply_test.cpp diff --git a/stan/math/fwd/core/eigen_product_hooks.hpp b/stan/math/fwd/core/eigen_product_hooks.hpp deleted file mode 100644 index 50b0a394980..00000000000 --- a/stan/math/fwd/core/eigen_product_hooks.hpp +++ /dev/null @@ -1,104 +0,0 @@ -#ifndef STAN_MATH_FWD_CORE_EIGEN_PRODUCT_HOOKS_HPP -#define STAN_MATH_FWD_CORE_EIGEN_PRODUCT_HOOKS_HPP - -#include -#include -#include -#include -#include -#include -#include - -/** - * @file - * Routes Eigen products whose operands have `fvar` or - * `std::complex>` scalars (mixed with arithmetic scalars) to the - * forward-mode kernels in stan/math/fwd/core/eigen_product_kernel.hpp, the - * same way stan/math/rev/core/eigen_product_hooks.hpp does for `var`. - */ - -namespace stan { -namespace math { -namespace internal { - -template -struct is_fwd_product_operand_scalar - : std::integral_constant::value> {}; -template -struct is_fwd_product_operand_scalar, T> - : std::integral_constant::value> {}; -template -struct is_fwd_product_operand_scalar, T> : std::true_type {}; -template -struct is_fwd_product_operand_scalar>, T> - : std::true_type {}; - -template -inline constexpr bool is_fwd_product_scalars_v - = !std::is_void::value && is_fwd_product_operand_scalar::value - && is_fwd_product_operand_scalar::value; - -/** - * The inner scalar `T` shared by the `fvar` operands of a product, or `void` - * if neither operand is an `fvar`. - */ -template -using fwd_product_inner_of_t - = std::conditional_t>::value, - fwd_product_inner_t, fwd_product_inner_t>; - -template -inline constexpr bool is_fwd_product_v = is_fwd_product_scalars_v< - typename Eigen::internal::traits::Scalar, - typename Eigen::internal::traits::Scalar, - fwd_product_inner_of_t::Scalar, - typename Eigen::internal::traits::Scalar>>; - -/** - * Forward-mode matrix product. - */ -template -struct eigen_matrix_product_impl>> { - using result_t = typename Eigen::Product::PlainObject; - static result_t run(const Lhs& lhs, const Rhs& rhs) { - return fwd_matrix_product(lhs, rhs); - } -}; - -/** - * Forward-mode coefficient-wise product, or product with a scalar. - */ -template -struct eigen_cwise_product_impl< - Eigen::CwiseBinaryOp, Lhs, Rhs>, - std::enable_if_t>> { - using xpr_t = Eigen::CwiseBinaryOp, - Lhs, Rhs>; - using result_t = typename xpr_t::PlainObject; - static auto run(const xpr_t& xpr) { return fwd_cwise_product(xpr); } -}; - -} // namespace internal -} // namespace math -} // namespace stan - -namespace Eigen { -namespace internal { - -#define STAN_FWD_PRODUCT_T_TPARAMS() typename T_, -#define STAN_FWD_PRODUCT_TS_TPARAMS() typename T_, typename S_, -STAN_EIGEN_AD_SCALAR_PATTERNS(STAN_EIGEN_AD_CWISE_PRODUCT_EVALUATOR, - STAN_FWD_PRODUCT_T_TPARAMS, - STAN_FWD_PRODUCT_TS_TPARAMS, stan::math::fvar) -STAN_EIGEN_AD_PRODUCT_ASSIGNMENT(STAN_FWD_PRODUCT_T_TPARAMS, - stan::math::fvar) -STAN_EIGEN_AD_PRODUCT_ASSIGNMENT(STAN_FWD_PRODUCT_T_TPARAMS, - std::complex>) -#undef STAN_FWD_PRODUCT_T_TPARAMS -#undef STAN_FWD_PRODUCT_TS_TPARAMS - -} // namespace internal -} // namespace Eigen - -#endif diff --git a/stan/math/fwd/core/eigen_product_kernel.hpp b/stan/math/fwd/core/eigen_product_kernel.hpp deleted file mode 100644 index 9a65ae2189e..00000000000 --- a/stan/math/fwd/core/eigen_product_kernel.hpp +++ /dev/null @@ -1,384 +0,0 @@ -#ifndef STAN_MATH_FWD_CORE_EIGEN_PRODUCT_KERNEL_HPP -#define STAN_MATH_FWD_CORE_EIGEN_PRODUCT_KERNEL_HPP - -#include -#include -#include -#include -#include -#include -#include -#include -#include - -/** - * @file - * Forward-mode kernels for Eigen products whose operands have `fvar` or - * `std::complex>` scalars. A product of `fvar` matrices is split - * into products of their values and tangents, which are Eigen products of - * the inner scalar `T`. For `T = var` those go through the reverse-mode - * kernels, so an `fvar` matrix product puts two callbacks on the - * reverse-mode stack instead of several varis per scalar operation; for - * `T = double` they are ordinary double products. - * - * The kernels must never form a coefficient-wise or matrix product with - * `fvar` scalars themselves, since that would dispatch back into these - * kernels. Lazy results use `fwd_lazy_product_op` instead of Eigen's - * `scalar_product_op`. - */ - -namespace stan { -namespace math { -namespace internal { - -/** - * The inner scalar type `T` of `fvar` and `std::complex>`, and - * `void` for other scalars. - */ -template -struct fwd_product_inner { - using type = void; -}; -template -struct fwd_product_inner> { - using type = T; -}; -template -struct fwd_product_inner>> { - using type = T; -}; -template -using fwd_product_inner_t = typename fwd_product_inner::type; - -/** - * True if the innermost scalar of `fvar<......>` is a reverse-mode - * autodiff type, so products at that level record reverse-mode nodes. - */ -template -struct fwd_product_uses_tape : is_var {}; -template -struct fwd_product_uses_tape> : fwd_product_uses_tape {}; -template -struct fwd_product_uses_tape> : fwd_product_uses_tape {}; - -/** - * The value type used for an operand of a product with inner scalar `T`: - * `T` itself for `fvar`, `std::complex` for `std::complex>`, - * and `double` or `std::complex` for arithmetic operands. - */ -template -struct fwd_product_value { - using type = double; -}; -template -struct fwd_product_value, - std::enable_if_t::value>> { - using type = std::complex; -}; -template -struct fwd_product_value> { - using type = T; -}; -template -struct fwd_product_value>> { - using type = std::complex; -}; - -template -inline const T& fwd_product_val(const fvar& x) { - return x.val_; -} -template -inline const T& fwd_product_tangent(const fvar& x) { - return x.d_; -} -template -inline std::complex fwd_product_val(const std::complex>& z) { - return {z.real().val_, z.imag().val_}; -} -template -inline std::complex fwd_product_tangent(const std::complex>& z) { - return {z.real().d_, z.imag().d_}; -} - -template -inline const T& fwd_product_real(const T& x) { - return x; -} -template -inline T fwd_product_real(const std::complex& z) { - return z.real(); -} -template -inline double fwd_product_imag(const T&) { - return 0.0; -} -template -inline T fwd_product_imag(const std::complex& z) { - return z.imag(); -} - -/** - * Build a result coefficient from its value and tangent, which may have a - * lower scalar type (e.g. `double` for products of arithmetic operands - * scaled by an `fvar`). - */ -template -inline void fwd_product_assign(fvar& res, const V& val, const D& d) { - res = fvar(T(val), T(d)); -} -template -inline void fwd_product_assign(std::complex>& res, const V& val, - const D& d) { - res = std::complex>( - fvar(T(fwd_product_real(val)), T(fwd_product_real(d))), - fvar(T(fwd_product_imag(val)), T(fwd_product_imag(d)))); -} - -template -using fwd_operand_plain_t - = Eigen::Matrix; - -/** - * Operand of a forward-mode product kernel: its values and, for autodiff - * scalars, its tangents, each a plain matrix of the inner scalar type. - * - * @tparam X operand expression type - */ -template ::Scalar, - bool HasTangent = !std::is_void>::value> -struct fwd_product_operand { - static constexpr bool has_tangent = false; - using value_t = typename fwd_product_value::type; - fwd_operand_plain_t val_; - explicit fwd_product_operand(const X& x) : val_(x.template cast()) {} -}; - -template -struct fwd_product_operand { - static constexpr bool has_tangent = true; - using value_t = typename fwd_product_value::type; - fwd_operand_plain_t val_; - fwd_operand_plain_t d_; - explicit fwd_product_operand(const X& x) { - const fwd_operand_plain_t x_eval = x; - val_ = x_eval.unaryExpr( - [](const Scalar& s) -> value_t { return fwd_product_val(s); }); - d_ = x_eval.unaryExpr( - [](const Scalar& s) -> value_t { return fwd_product_tangent(s); }); - } -}; - -/** - * Scalar factors folded into a forward-mode product kernel, multiplied - * together as a scalar of the result type. - * - * @tparam ResScalar scalar type of the product result - */ -template -struct fwd_product_scales { - // empty until the first factor: constructing an `fvar`, even a zero - // or one, puts varis on the reverse-mode stack - std::optional value_; - bool any() const { return value_.has_value(); } - template - void push(const S& s) { - if (value_) { - *value_ *= s; - } else { - value_.emplace(s); - } - } -}; - -/** - * Write the value and tangent matrices of a kernel result into its plain - * result type. - */ -template -inline Result fwd_product_make_result(const Val& val, const Tangent& d) { - Result res(val.rows(), val.cols()); - for (Eigen::Index j = 0; j < val.cols(); ++j) { - for (Eigen::Index i = 0; i < val.rows(); ++i) { - fwd_product_assign(res.coeffRef(i, j), val.coeff(i, j), d.coeff(i, j)); - } - } - return res; -} - -/** - * The value and tangent of `sigma * M` from those of `M`, where `sigma` is - * the product of the folded scalar factors. - */ -template -inline Result fwd_product_scaled_result(const Scales& scales, const MVal& m_val, - const MTangent& m_d) { - if (!scales.any()) { - return fwd_product_make_result(m_val, m_d); - } - const auto s_val = fwd_product_val(*scales.value_); - const auto s_d = fwd_product_tangent(*scales.value_); - const auto val = (s_val * m_val).eval(); - const auto d = (s_d * m_val + s_val * m_d).eval(); - return fwd_product_make_result(val, d); -} - -/** - * `sigma * X * Y` for matrix operands. The tangent is `X.val * Y.d + X.d * - * Y.val`, computed as one product of the concatenations `[X.val, X.d]` and - * `[Y.d; Y.val]` when both operands have tangents. - */ -template -inline Result fwd_matmul(const Scales& scales, const XOp& x, const YOp& y) { - const auto m_val = (x.val_ * y.val_).eval(); - if constexpr (XOp::has_tangent && YOp::has_tangent) { - using x_value_t = typename XOp::value_t; - using y_value_t = typename YOp::value_t; - Eigen::Matrix x_cat( - x.val_.rows(), 2 * x.val_.cols()); - x_cat << x.val_, x.d_; - Eigen::Matrix y_cat( - 2 * y.val_.rows(), y.val_.cols()); - y_cat << y.d_, y.val_; - const auto m_d = (x_cat * y_cat).eval(); - return fwd_product_scaled_result(scales, m_val, m_d); - } else if constexpr (XOp::has_tangent) { - const auto m_d = (x.d_ * y.val_).eval(); - return fwd_product_scaled_result(scales, m_val, m_d); - } else if constexpr (YOp::has_tangent) { - const auto m_d = (x.val_ * y.d_).eval(); - return fwd_product_scaled_result(scales, m_val, m_d); - } else { - // arithmetic operands: the tangent comes from the scalar factors only - return fwd_product_scaled_result( - scales, m_val, - std::decay_t::Zero(m_val.rows(), m_val.cols())); - } -} - -/** - * `sigma * (X .* Y)` for operands of equal size. - */ -template -inline Result fwd_cwise(const Scales& scales, const XOp& x, const YOp& y) { - const auto m_val = (x.val_.array() * y.val_.array()).matrix().eval(); - if constexpr (XOp::has_tangent && YOp::has_tangent) { - const auto m_d - = (x.val_.array() * y.d_.array() + x.d_.array() * y.val_.array()) - .matrix() - .eval(); - return fwd_product_scaled_result(scales, m_val, m_d); - } else if constexpr (XOp::has_tangent) { - const auto m_d = (x.d_.array() * y.val_.array()).matrix().eval(); - return fwd_product_scaled_result(scales, m_val, m_d); - } else if constexpr (YOp::has_tangent) { - const auto m_d = (x.val_.array() * y.d_.array()).matrix().eval(); - return fwd_product_scaled_result(scales, m_val, m_d); - } else { - // arithmetic operands: the tangent comes from the scalar factors only - return fwd_product_scaled_result( - scales, m_val, - std::decay_t::Zero(m_val.rows(), m_val.cols())); - } -} - -/** - * `sigma * X` for a matrix operand. - */ -template -inline Result fwd_scale(const Scales& scales, const XOp& x) { - if constexpr (XOp::has_tangent) { - return fwd_product_scaled_result(scales, x.val_, x.d_); - } else { - return fwd_product_scaled_result( - scales, x.val_, - std::decay_t::Zero(x.val_.rows(), x.val_.cols()) - .eval()); - } -} - -/** - * Matrix product `lhs * rhs` of Eigen expressions with `fvar` scalars. - * - * @tparam Result plain Eigen type of the product - */ -template -inline Result fwd_matrix_product(const Lhs& lhs, const Rhs& rhs) { - fwd_product_scales scales; - const auto& lhs_core = peel_scales(scales, lhs); - const auto& rhs_core = peel_scales(scales, rhs); - using lhs_core_t = std::decay_t; - using rhs_core_t = std::decay_t; - return fwd_matmul(scales, fwd_product_operand(lhs_core), - fwd_product_operand(rhs_core)); -} - -/** - * Coefficient-wise product functor for lazy forward-mode products. It is - * not Eigen's `scalar_product_op`, so expressions using it are evaluated by - * Eigen's generic evaluator instead of Stan's product hooks. - */ -template -struct fwd_lazy_product_op { - using result_type = Result; - Result operator()(const L& a, const R& b) const { return a * b; } -}; - -/** - * Coefficient-wise product expression `xpr` (including a matrix times a - * scalar) with `fvar` scalars. - * - * When no reverse-mode tape is involved (e.g. `fvar`), splitting - * into value and tangent products would only add temporaries, so the - * product stays a lazy per-coefficient expression. Otherwise the value and - * tangent products are evaluated as products of the inner scalar, and a - * scaled matrix product `(A * B) * s` becomes one matrix-product kernel. - * - * @tparam Result plain Eigen type of the expression - */ -template -inline auto fwd_cwise_product(const Xpr& xpr) { - using res_scalar = typename Result::Scalar; - if constexpr (!fwd_product_uses_tape::value) { - using lhs_scalar = typename Eigen::internal::traits< - std::decay_t>::Scalar; - using rhs_scalar = typename Eigen::internal::traits< - std::decay_t>::Scalar; - return xpr.lhs().binaryExpr( - xpr.rhs(), fwd_lazy_product_op()); - } else { - fwd_product_scales scales; - const auto& core = peel_scales(scales, xpr); - using core_t = std::decay_t; - if constexpr (std::is_same::value) { - const auto& lhs_core = peel_scales(scales, xpr.lhs()); - const auto& rhs_core = peel_scales(scales, xpr.rhs()); - using lhs_core_t = std::decay_t; - using rhs_core_t = std::decay_t; - return fwd_cwise(scales, - fwd_product_operand(lhs_core), - fwd_product_operand(rhs_core)); - } else if constexpr (is_dense_product::value) { - const auto& lhs_core = peel_scales(scales, core.lhs()); - const auto& rhs_core = peel_scales(scales, core.rhs()); - using lhs_core_t = std::decay_t; - using rhs_core_t = std::decay_t; - return fwd_matmul(scales, - fwd_product_operand(lhs_core), - fwd_product_operand(rhs_core)); - } else { - return fwd_scale(scales, fwd_product_operand(core)); - } - } -} - -} // namespace internal -} // namespace math -} // namespace stan - -#endif diff --git a/stan/math/fwd/core/fvar.hpp b/stan/math/fwd/core/fvar.hpp index cca3d751f87..187d6a5700b 100644 --- a/stan/math/fwd/core/fvar.hpp +++ b/stan/math/fwd/core/fvar.hpp @@ -1,7 +1,6 @@ #ifndef STAN_MATH_FWD_CORE_FVAR_HPP #define STAN_MATH_FWD_CORE_FVAR_HPP -#include #include #include #include @@ -257,16 +256,6 @@ struct fvar { } }; -namespace internal { -/** - * Eigen products of `fvar` scalars are evaluated by Stan (see - * stan/math/prim/meta/eigen_product_evaluator_hooks.hpp). Declared with - * `fvar` itself so every translation unit that uses `fvar` agrees. - */ -template -struct is_stan_ad_product_scalar> : std::true_type {}; -} // namespace internal - } // namespace math } // namespace stan #endif diff --git a/stan/math/fwd/fun/Eigen_NumTraits.hpp b/stan/math/fwd/fun/Eigen_NumTraits.hpp index 9816fcc6d0d..d4577f20d4c 100644 --- a/stan/math/fwd/fun/Eigen_NumTraits.hpp +++ b/stan/math/fwd/fun/Eigen_NumTraits.hpp @@ -253,6 +253,5 @@ struct ternary_evaluator< } // namespace internal } // namespace Eigen -#include #endif diff --git a/stan/math/prim/fun.hpp b/stan/math/prim/fun.hpp index a09499863d4..7d0dbfa53cf 100644 --- a/stan/math/prim/fun.hpp +++ b/stan/math/prim/fun.hpp @@ -223,6 +223,7 @@ #include #include #include +#include #include #include #include @@ -281,6 +282,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..8dddd5003db --- /dev/null +++ b/stan/math/prim/fun/multiply_add.hpp @@ -0,0 +1,41 @@ +#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(const Mat1& A, const Mat2& B, const 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 (A * B + C).eval(); +} + +} // 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..1a79d9e8592 --- /dev/null +++ b/stan/math/prim/fun/scale_multiply.hpp @@ -0,0 +1,36 @@ +#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(const Scal& s, const Mat1& A, const Mat2& B) { + check_multiplicable("scale_multiply", "A", A, "B", B); + return (s * (A * B)).eval(); +} + +} // namespace math +} // namespace stan + +#endif diff --git a/stan/math/prim/meta/eigen_product_evaluator_hooks.hpp b/stan/math/prim/meta/eigen_product_evaluator_hooks.hpp index 707bb40fcec..117bc61d802 100644 --- a/stan/math/prim/meta/eigen_product_evaluator_hooks.hpp +++ b/stan/math/prim/meta/eigen_product_evaluator_hooks.hpp @@ -3,394 +3,94 @@ /** * @file - * Hooks that let Stan evaluate Eigen matrix products whose operands have - * autodiff scalars. + * Specializations of Eigen's evaluators that compute products of matrices + * with autodiff scalars with Stan's `multiply`, `elt_multiply`, + * `scale_multiply` and `multiply_add`, so that `A * B` records one node on + * the reverse-mode stack instead of one per scalar operation. * - * Eigen dispatches `lhs * rhs` for dense operands through - * `Eigen::internal::product_evaluator`. Its scalar template parameters are - * defaulted from the operand types, so a partial specialization that fixes - * them is ambiguous with Eigen's generic specialization. Instead, the - * specializations below match every dense-by-dense product and choose, from - * the operand scalar types, between + * Eigen's `product_evaluator` takes the operand scalar types as defaulted + * template parameters, so a specialization that fixes them is ambiguous with + * Eigen's generic one. The dense `product_evaluator` below therefore matches + * every dense-by-dense product and handles non-autodiff scalars exactly as + * Eigen's generic `product_evaluator` in Eigen/src/Core/ProductEvaluators.h + * does. When upgrading Eigen, compare that evaluator with this one. * - * - `stan::math::internal::ad_product_evaluator`, which forwards to - * `eigen_matrix_product_impl::run(lhs, rhs)`, a customization - * point defined by the autodiff layers (`rev`, `fwd`); and - * - Eigen's own evaluation for all other scalars (e.g. `double`), which is - * reproduced exactly: `eigen_dense_product_evaluator` is a copy of the - * generic `product_evaluator` in Eigen/src/Core/ProductEvaluators.h, and - * lazy products derive from Eigen's coefficient-based evaluator. - * - * Because these specializations are selected for every dense product, they - * are included from stan/math/prim/fun/Eigen.hpp so that every translation - * unit sees the same definition. When upgrading Eigen, check the generic - * `product_evaluator` and the `LazyProduct` evaluator in - * Eigen/src/Core/ProductEvaluators.h against the copies here. + * Because the specializations match every dense product, this file is + * included from stan/math/prim/fun/Eigen.hpp so that every translation unit + * sees them. The functions they 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 namespace stan { namespace math { namespace internal { /** - * True for scalar types whose Eigen products Stan evaluates itself. The - * primary template is false; `rev` specializes it for `var` and `fwd` for - * `fvar`. Complex scalars inherit from their value type. - * - * @tparam Scalar scalar type of an Eigen expression - */ -template -struct is_stan_ad_product_scalar : std::false_type {}; - -template -struct is_stan_ad_product_scalar> - : is_stan_ad_product_scalar {}; - -/** - * True if either operand of an Eigen product has a scalar type for which - * Stan evaluates the product. + * 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_stan_ad_product_scalar< - typename Eigen::internal::traits::Scalar>::value - || is_stan_ad_product_scalar< - typename Eigen::internal::traits::Scalar>::value; - -/** - * Customization point for dense matrix products with autodiff operands. - * Specializations provide `static auto run(const Lhs&, const Rhs&)`, which - * returns an object with an Eigen evaluator holding the product, usually the - * product's plain object type. The primary template is intentionally left - * undefined so that a product with an autodiff scalar and no implementation - * is a compile error. - * - * @tparam Lhs left operand expression type - * @tparam Rhs right operand expression type - * @tparam Enable used by specializations for SFINAE - */ -template -struct eigen_matrix_product_impl; - -/** - * Customization point for coefficient-wise products (including the product - * of a matrix and a scalar) with autodiff operands. Specializations provide - * `static auto run(const Xpr&)` for the `CwiseBinaryOp` expression `Xpr`, - * returning either an evaluated plain object or a lazy expression; the - * implementation decides which. - * - * @tparam Xpr the `CwiseBinaryOp` of a `scalar_product_op` - * @tparam Enable used by specializations for SFINAE - */ -template -struct eigen_cwise_product_impl; - -/** - * Owns the object returned by a product customization point. It is a base - * class of `forwarding_evaluator` listed before the evaluator base, so the - * result is constructed before the evaluator that refers to it. - * - * @tparam Result type returned by the customization point - */ -template -struct forwarded_result { - Result result_; - template - explicit forwarded_result(F&& make) : result_(make()) {} -}; + = (is_var::Scalar>::value + || is_var::Scalar>::value) + && !is_complex::Scalar>::value + && !is_complex::Scalar>::value; /** - * Eigen evaluator for an expression whose value is computed by a Stan - * customization point. The returned object is stored and read through its - * own Eigen evaluator. When the returned object is a plain matrix or array - * the evaluator sets `EvalBeforeNestingBit`, so an enclosing expression - * evaluates this one once instead of once per coefficient access. - * - * Eigen never copies evaluators (`evaluator_base` is noncopyable), so the - * inner evaluator's pointer into `result_` stays valid. + * 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 Result type returned by the customization point - */ -template -struct forwarding_evaluator : private forwarded_result, - public Eigen::internal::evaluator { - using holder_t = forwarded_result; - using Base = Eigen::internal::evaluator; - static constexpr bool is_plain_result - = std::is_base_of, Result>::value; - enum { - Flags - = static_cast(Base::Flags) - | (is_plain_result ? static_cast(Eigen::EvalBeforeNestingBit) : 0) - }; - - template - explicit forwarding_evaluator(F&& make) - : holder_t(std::forward(make)), Base(this->holder_t::result_) {} -}; - -template -using matrix_product_result_t - = std::decay_t::run( - std::declval(), std::declval()))>; - -/** - * Evaluator of a dense product with an autodiff operand. It forwards to - * `eigen_matrix_product_impl`. - * - * @tparam Lhs left operand expression type - * @tparam Rhs right operand expression type - * @tparam Options Eigen product options (default, lazy, alias free) - */ -template -struct ad_product_evaluator - : forwarding_evaluator> { - using XprType = Eigen::Product; - using Base = forwarding_evaluator>; - static constexpr bool is_stan_ad_product_hook = true; - explicit ad_product_evaluator(const XprType& xpr) - : Base([&xpr]() { - return eigen_matrix_product_impl::run(xpr.lhs(), xpr.rhs()); - }) {} -}; - -/** - * Evaluator of a coefficient-wise product with an autodiff operand. It - * forwards to `eigen_cwise_product_impl`. - * - * @tparam Xpr the `CwiseBinaryOp` of a `scalar_product_op` - */ -template -struct ad_cwise_product_evaluator - : forwarding_evaluator::run(std::declval()))>> { - using XprType = Xpr; - using Base = forwarding_evaluator::run(std::declval()))>>; - static constexpr bool is_stan_ad_product_hook = true; - explicit ad_cwise_product_evaluator(const Xpr& xpr) - : Base([&xpr]() { return eigen_cwise_product_impl::run(xpr); }) {} -}; - -/** - * Detects the scalar-times-expression shapes Eigen builds for `s * X` and - * `X * s` (see `EIGEN_SCALAR_BINARYOP_EXPR_RETURN_TYPE` in Eigen's - * util/Macros.h), so a kernel can fold the scalar into one callback. + * @tparam T an Eigen expression type */ template -struct scaled_expression : std::false_type {}; - -template -struct scaled_expression +struct is_scalar_times_matrix, const Eigen::CwiseNullaryOp, Plain>, - const Inner>> : std::true_type { - template - static const S1& scale(const X& x) { - return x.lhs().functor().m_other; - } - template - static const Inner& inner(const X& x) { - return x.rhs(); - } -}; - -template -struct scaled_expression, const Inner, - const Eigen::CwiseNullaryOp, - Plain>>> : std::true_type { - template - static const S2& scale(const X& x) { - return x.rhs().functor().m_other; - } - template - static const Inner& inner(const X& x) { - return x.lhs(); - } -}; - -template -struct scaled_expression< - Eigen::CwiseBinaryOp, - const Eigen::CwiseNullaryOp< - Eigen::internal::scalar_constant_op, Plain1>, - const Eigen::CwiseNullaryOp< - Eigen::internal::scalar_constant_op, Plain2>>> - : std::true_type { - using Inner - = const Eigen::CwiseNullaryOp, - Plain2>; - template - static const S1& scale(const X& x) { - return x.lhs().functor().m_other; - } - template - static Inner& inner(const X& x) { - return x.rhs(); - } -}; + const Mat>> : std::true_type {}; /** - * Maximum number of nested scalar factors folded out of one operand. Deeper - * nests are left in the operand, which then goes through its own product - * kernel. - */ -constexpr int max_peel_depth = 4; - -/** - * Strip up to `max_peel_depth` scalar factors off `x` into `scales` and - * return the remaining expression. - */ -template -inline decltype(auto) peel_scales(Scales& scales, const X& x) { - if constexpr (scaled_expression::value && Depth < max_peel_depth) { - scales.push(scaled_expression::scale(x)); - return peel_scales(scales, scaled_expression::inner(x)); - } else { - return static_cast(x); - } -} - -/** - * True for an Eigen product of two dense operands. - */ -template -struct is_dense_product : std::false_type {}; -template -struct is_dense_product> - : std::integral_constant< - bool, - std::is_same::Shape, - Eigen::DenseShape>::value - && std::is_same< - typename Eigen::internal::evaluator_traits::Shape, - Eigen::DenseShape>::value> {}; - -/** - * True if the Eigen evaluator `Evaluator` is one of Stan's autodiff product - * evaluators, i.e. the expression it evaluates is computed by Stan. - * - * @tparam Evaluator an `Eigen::internal::evaluator` type - */ -template -struct is_stan_ad_product_hook : std::false_type {}; -template -struct is_stan_ad_product_hook< - Evaluator, std::void_t> - : std::integral_constant {}; - -/** - * True if Eigen evaluates the expression `Xpr` with one of Stan's autodiff - * product evaluators. + * 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 Xpr an Eigen expression type + * @tparam T an Eigen expression type */ -template -inline constexpr bool is_stan_ad_product_hook_v - = is_stan_ad_product_hook>::value; - -/** - * Eigen's generic product evaluator for dense operands, reproduced from - * Eigen/src/Core/ProductEvaluators.h (`product_evaluator, ProductTag, LhsShape, RhsShape>`). It is used for products of - * scalars that Stan does not evaluate itself. - * - * @tparam Lhs left operand expression type - * @tparam Rhs right operand expression type - * @tparam Options Eigen product options - * @tparam ProductTag Eigen product implementation type - */ -template -struct eigen_dense_product_evaluator - : public Eigen::internal::evaluator< - typename Eigen::Product::PlainObject> { - using XprType = Eigen::Product; - using PlainObject = typename XprType::PlainObject; - using Base = Eigen::internal::evaluator; - enum { Flags = Base::Flags | Eigen::EvalBeforeNestingBit }; - - EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE explicit eigen_dense_product_evaluator( - const XprType& xpr) - : m_result(xpr.rows(), xpr.cols()) { - Eigen::internal::construct_at(this, m_result); - Eigen::internal::generic_product_impl::evalTo(m_result, - xpr.lhs(), - xpr.rhs()); - } - - protected: - PlainObject m_result; -}; +template +struct is_matrix_times_scalar : std::false_type {}; +template +struct is_matrix_times_scalar< + Eigen::CwiseBinaryOp, const Mat, + const Eigen::CwiseNullaryOp< + Eigen::internal::scalar_constant_op, Plain>>> + : std::true_type {}; /** - * `Assignment` of a product to a dense destination with the same autodiff - * scalar. Eigen's own specializations for this case call - * `generic_product_impl` directly and skip `product_evaluator`; for dense - * operands this forwards to `eigen_matrix_product_impl` instead. Products - * with non-dense operands (e.g. triangular views) run Eigen's body. + * True for an Eigen matrix product expression. * - * @tparam Dst destination expression type - * @tparam Lhs left operand expression type - * @tparam Rhs right operand expression type - * @tparam Options Eigen product options - * @tparam Func assignment functor (`assign_op`, `add_assign_op`, - * `sub_assign_op`) + * @tparam T an Eigen expression type */ -template -struct ad_product_assignment { - using SrcXprType = Eigen::Product; - using Scalar = typename Dst::Scalar; - static constexpr bool dense_operands - = std::is_same::Shape, - Eigen::DenseShape>::value - && std::is_same::Shape, - Eigen::DenseShape>::value; - static void run(Dst& dst, const SrcXprType& src, const Func& func) { - if constexpr (dense_operands) { - const auto result - = eigen_matrix_product_impl::run(src.lhs(), src.rhs()); - Eigen::internal::call_assignment_no_alias(dst, result, func); - } else if constexpr (std::is_same>::value) { - if (dst.rows() != src.rows() || dst.cols() != src.cols()) { - dst.resize(src.rows(), src.cols()); - } - Eigen::internal::generic_product_impl::evalTo(dst, src.lhs(), - src.rhs()); - } else if constexpr (std::is_same>::value) { - Eigen::internal::generic_product_impl::addTo(dst, src.lhs(), - src.rhs()); - } else { - Eigen::internal::generic_product_impl::subTo(dst, src.lhs(), - src.rhs()); - } - } -}; - -/** - * Tag passed to Eigen's lazy (coefficient-based) product evaluator when it is - * reached from one of Stan's per-tag specializations. Eigen's lazy evaluator - * does not depend on its product tag, and no Eigen product has this tag, so - * it resolves to Eigen's own evaluator. - */ -constexpr int eigen_lazy_product_passthrough_tag = -1; - -template -using eigen_lazy_product_evaluator = Eigen::internal::product_evaluator< - Eigen::Product, - eigen_lazy_product_passthrough_tag, Eigen::DenseShape, Eigen::DenseShape>; +template +struct is_eigen_product : std::false_type {}; +template +struct is_eigen_product> : std::true_type {}; } // namespace internal } // namespace math @@ -400,46 +100,99 @@ namespace Eigen { namespace internal { /** - * Dense-by-dense products: Stan's evaluator for autodiff operands, Eigen's - * generic evaluator otherwise. More specialized than Eigen's generic - * `product_evaluator` (fixed shapes) and less specialized than Eigen's lazy - * and `LazyCoeffBasedProductMode` evaluators, which are handled below. + * Dense-by-dense products. For operands with `var` scalars the product is + * `scale_multiply(s, A, B)` when either operand is Eigen's `s * A` (Eigen + * rewrites `s * (A * B)` as `(s * A) * B`) and `multiply(lhs, rhs)` + * otherwise. For all other scalars this is Eigen's generic + * `product_evaluator`. More specialized than Eigen's generic + * `product_evaluator` (fixed shapes) and less specialized than its lazy and + * `LazyCoeffBasedProductMode` evaluators, which are handled below. */ template struct product_evaluator, ProductTag, DenseShape, DenseShape> - : std::conditional_t< - stan::math::internal::is_stan_ad_product_v, - stan::math::internal::ad_product_evaluator, - stan::math::internal::eigen_dense_product_evaluator> { - using Base = std::conditional_t< - stan::math::internal::is_stan_ad_product_v, - stan::math::internal::ad_product_evaluator, - stan::math::internal::eigen_dense_product_evaluator>; - using Base::Base; + : public evaluator::PlainObject> { + using XprType = Product; + using PlainObject = typename XprType::PlainObject; + using Base = evaluator; + enum { Flags = Base::Flags | EvalBeforeNestingBit }; + + explicit product_evaluator(const XprType& xpr) + : m_result(xpr.rows(), xpr.cols()) { + using stan::math::internal::is_matrix_times_scalar; + using stan::math::internal::is_scalar_times_matrix; + if constexpr (!stan::math::internal::is_stan_ad_product_v) { + generic_product_impl::evalTo(m_result, xpr.lhs(), xpr.rhs()); + } else if constexpr (is_scalar_times_matrix::value) { + // (s * A) * B + m_result = scale_multiply(xpr.lhs().lhs().functor().m_other, + xpr.lhs().rhs(), xpr.rhs()); + } else if constexpr (is_matrix_times_scalar::value) { + // (A * s) * B + m_result = scale_multiply(xpr.lhs().rhs().functor().m_other, + xpr.lhs().lhs(), xpr.rhs()); + } else if constexpr (is_scalar_times_matrix::value) { + // A * (s * B) + m_result = scale_multiply(xpr.rhs().lhs().functor().m_other, xpr.lhs(), + xpr.rhs().rhs()); + } else if constexpr (is_matrix_times_scalar::value) { + // A * (B * s) + m_result = scale_multiply(xpr.rhs().rhs().functor().m_other, xpr.lhs(), + xpr.rhs().lhs()); + } else if constexpr (stan::is_stan_scalar::value) { + // a row vector times a column vector + m_result.coeffRef(0, 0) = multiply(xpr.lhs(), xpr.rhs()); + } else { + m_result = multiply(xpr.lhs(), xpr.rhs()); + } + // 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. + construct_at(this, m_result); + } + + protected: + PlainObject m_result; }; /** * Lazy (coefficient-based) products, one specialization per product tag so - * each is more specialized than Eigen's `LazyProduct` evaluator. Eigen's + * each is more specialized than Eigen's `LazyProduct` evaluator. For `var` + * operands they use the dense evaluator above (the product tag is unused + * there); for other scalars Eigen's lazy evaluator, reached with tag -1, + * which no product has and which Eigen's lazy evaluator ignores. Eigen's * `LazyCoeffBasedProductMode` evaluator for default products derives from - * the `CoeffBasedProductMode` one, so it is covered as well. + * the `CoeffBasedProductMode` specialization, so it is covered as well. */ -#define STAN_EIGEN_LAZY_PRODUCT_EVALUATOR(TAG) \ - template \ - struct product_evaluator, TAG, DenseShape, \ - DenseShape> \ - : std::conditional_t< \ - stan::math::internal::is_stan_ad_product_v, \ - stan::math::internal::ad_product_evaluator, \ - stan::math::internal::eigen_lazy_product_evaluator> { \ - using Base = std::conditional_t< \ - stan::math::internal::is_stan_ad_product_v, \ - stan::math::internal::ad_product_evaluator, \ - stan::math::internal::eigen_lazy_product_evaluator>; \ - using Base::Base; \ +#define STAN_EIGEN_LAZY_PRODUCT_EVALUATOR(TAG) \ + template \ + struct product_evaluator, TAG, DenseShape, \ + DenseShape> \ + : std::conditional_t< \ + stan::math::internal::is_stan_ad_product_v, \ + product_evaluator, GemmProduct, \ + DenseShape, DenseShape>, \ + product_evaluator, -1, DenseShape, \ + DenseShape>> { \ + using XprType = Product; \ + using Base = std::conditional_t< \ + stan::math::internal::is_stan_ad_product_v, \ + product_evaluator, GemmProduct, \ + DenseShape, DenseShape>, \ + product_evaluator, -1, DenseShape, \ + DenseShape>>; \ + explicit product_evaluator(const XprType& xpr) : Base(base_xpr(xpr)) {} \ + \ + private: \ + static auto base_xpr(const XprType& xpr) { \ + if constexpr (stan::math::internal::is_stan_ad_product_v) { \ + return Product(xpr.lhs(), xpr.rhs()); \ + } else { \ + return xpr; \ + } \ + } \ }; STAN_EIGEN_LAZY_PRODUCT_EVALUATOR(InnerProduct) @@ -454,58 +207,168 @@ STAN_EIGEN_LAZY_PRODUCT_EVALUATOR(GemmProduct) } // namespace Eigen /** - * Apply `X(TPL, LhsScalar, RhsScalar)` to every scalar pattern of a product - * with an autodiff operand of type `AD` (and `std::complex`), where `S_` - * stands for any other scalar. `TPL` and `TPL_S` are function-like macros - * expanding to the extra template parameters (with a trailing comma) of the - * patterns without and with `S_`. The `(AD, AD)`, `(complex, - * complex)`, `(AD, complex)` and `(complex, AD)` patterns break - * the ties between the patterns with `S_`. - */ -#define STAN_EIGEN_AD_SCALAR_PATTERNS(X, TPL, TPL_S, AD) \ - X(TPL_S, AD, S_) \ - X(TPL_S, S_, AD) \ - X(TPL, AD, AD) \ - X(TPL_S, std::complex, S_) \ - X(TPL_S, S_, std::complex) \ - X(TPL, std::complex, std::complex) \ - X(TPL, AD, std::complex) \ - X(TPL, std::complex, AD) - -/** - * Specialize Eigen's `binary_evaluator` for coefficient-wise products with - * scalar pattern `(LS, RS)`, forwarding to `eigen_cwise_product_impl`. The - * scalar template parameters of `binary_evaluator` stay defaulted, exactly as - * in Eigen's generic specialization, so this one is more specialized only - * through the functor type. Use inside `namespace Eigen::internal`. + * Specialize Eigen's `binary_evaluator` for the coefficient-wise product + * `lhs .* rhs` with scalar pattern `(LS, RS)`. The product is + * `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. `TPL` is a + * function-like macro expanding to extra template parameters with a + * trailing comma. The scalar template parameters of `binary_evaluator` stay + * defaulted, as in Eigen's generic specialization, so this one is more + * specialized only through the functor. Use inside `namespace + * Eigen::internal`. */ #define STAN_EIGEN_AD_CWISE_PRODUCT_EVALUATOR(TPL, LS, RS) \ template \ struct binary_evaluator, Lhs, Rhs>, \ IndexBased, IndexBased> \ - : stan::math::internal::ad_cwise_product_evaluator< \ - CwiseBinaryOp, Lhs, Rhs>> { \ - using Base = stan::math::internal::ad_cwise_product_evaluator< \ - CwiseBinaryOp, Lhs, Rhs>>; \ - using Base::Base; \ + : public evaluator, \ + Lhs, Rhs>::PlainObject> { \ + using XprType = CwiseBinaryOp, Lhs, Rhs>; \ + using PlainObject = typename XprType::PlainObject; \ + using Base = evaluator; \ + enum { Flags = Base::Flags | EvalBeforeNestingBit }; \ + \ + explicit binary_evaluator(const XprType& xpr) \ + : m_result(xpr.rows(), xpr.cols()) { \ + using stan::math::internal::is_eigen_product; \ + using stan::math::internal::is_matrix_times_scalar; \ + using stan::math::internal::is_scalar_times_matrix; \ + /* Stan's functions take matrices, so array operands are viewed as */ \ + /* matrices (`.matrix().derived()` is the operand itself for a */ \ + /* matrix) and the result is viewed back as an array */ \ + const auto product = [&xpr]() { \ + if constexpr (is_scalar_times_matrix::value) { \ + if constexpr (is_eigen_product>::value) { \ + /* s * (A * B) */ \ + return scale_multiply(xpr.lhs().functor().m_other, \ + xpr.rhs().lhs(), xpr.rhs().rhs()); \ + } else { \ + /* s * X */ \ + return multiply(xpr.lhs().functor().m_other, \ + xpr.rhs().matrix().derived()); \ + } \ + } else if constexpr (is_matrix_times_scalar::value) { \ + if constexpr (is_eigen_product>::value) { \ + /* (A * B) * s */ \ + return scale_multiply(xpr.rhs().functor().m_other, \ + xpr.lhs().lhs(), xpr.lhs().rhs()); \ + } else { \ + /* X * s */ \ + return multiply(xpr.rhs().functor().m_other, \ + xpr.lhs().matrix().derived()); \ + } \ + } else { \ + return elt_multiply(xpr.lhs().matrix().derived(), \ + xpr.rhs().matrix().derived()); \ + } \ + }(); \ + if constexpr (std::is_base_of, \ + PlainObject>::value) { \ + m_result = product.array(); \ + } else { \ + m_result = product; \ + } \ + /* Base points into m_result's storage, so it is constructed only */ \ + /* once m_result is filled */ \ + construct_at(this, m_result); \ + } \ + \ + protected: \ + PlainObject m_result; \ }; /** - * Specialize Eigen's `Assignment` of a product to a dense destination with - * scalar `SCALAR` for `OP` (`assign_op`, `add_assign_op`, `sub_assign_op`). + * Specialize Eigen's `Assignment` of a product to a destination with scalar + * `SCALAR` for `OP` (`assign_op`, `add_assign_op` or `sub_assign_op`). + * Eigen's own specializations for this case compute the product directly + * instead of through `product_evaluator`. For dense operands, `dst = A * B` + * uses `multiply` (or `scale_multiply` for a scaled operand), `dst += A * B` + * uses `multiply_add(A, B, dst)` and `dst -= A * B` uses `subtract(dst, + * multiply(A, B))`; other shapes (e.g. triangular views) run Eigen's body. * The last template argument repeats Eigen's own condition token for token, * so this is more specialized only through the functor's scalar. Use inside * `namespace Eigen::internal`. */ -#define STAN_EIGEN_AD_PRODUCT_ASSIGNMENT_OP(TPL, SCALAR, OP) \ - template \ - struct Assignment, \ - OP, Dense2Dense, \ - std::enable_if_t<(Options == DefaultProduct \ - || Options == AliasFreeProduct)>> \ - : stan::math::internal::ad_product_assignment< \ - DstXprType, Lhs, Rhs, Options, OP> {}; +#define STAN_EIGEN_AD_PRODUCT_ASSIGNMENT_OP(TPL, SCALAR, OP) \ + template \ + struct Assignment, \ + OP, Dense2Dense, \ + std::enable_if_t<(Options == DefaultProduct \ + || Options == AliasFreeProduct)>> { \ + using SrcXprType = Product; \ + static void run(DstXprType& dst, const SrcXprType& src, \ + const OP& func) { \ + using stan::math::internal::is_matrix_times_scalar; \ + using stan::math::internal::is_scalar_times_matrix; \ + constexpr bool is_add \ + = std::is_same, \ + add_assign_op>::value; \ + constexpr bool is_sub \ + = std::is_same, \ + sub_assign_op>::value; \ + constexpr bool is_scaled = is_scalar_times_matrix::value \ + || is_matrix_times_scalar::value \ + || is_scalar_times_matrix::value \ + || is_matrix_times_scalar::value; \ + if constexpr (!std::is_same::Shape, \ + DenseShape>::value \ + || !std::is_same::Shape, \ + DenseShape>::value) { \ + if constexpr (is_add) { \ + generic_product_impl::addTo(dst, src.lhs(), src.rhs()); \ + } else if constexpr (is_sub) { \ + generic_product_impl::subTo(dst, src.lhs(), src.rhs()); \ + } else { \ + if (dst.rows() != src.rows() || dst.cols() != src.cols()) { \ + dst.resize(src.rows(), src.cols()); \ + } \ + generic_product_impl::evalTo(dst, src.lhs(), src.rhs()); \ + } \ + } else 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); \ + } \ + func.assignCoeff(dst.coeffRef(0, 0), multiply(src.lhs(), src.rhs())); \ + } else if constexpr (is_add && !is_scaled) { \ + dst = multiply_add(src.lhs(), src.rhs(), dst); \ + } else if constexpr (is_sub && !is_scaled) { \ + dst = subtract(dst, multiply(src.lhs(), src.rhs())); \ + } else { \ + const auto product = [&src]() { \ + if constexpr (is_scalar_times_matrix::value) { \ + /* (s * A) * B */ \ + return scale_multiply(src.lhs().lhs().functor().m_other, \ + src.lhs().rhs(), src.rhs()); \ + } else if constexpr (is_matrix_times_scalar::value) { \ + /* (A * s) * B */ \ + return scale_multiply(src.lhs().rhs().functor().m_other, \ + src.lhs().lhs(), src.rhs()); \ + } else if constexpr (is_scalar_times_matrix::value) { \ + /* A * (s * B) */ \ + return scale_multiply(src.rhs().lhs().functor().m_other, \ + src.lhs(), src.rhs().rhs()); \ + } else if constexpr (is_matrix_times_scalar::value) { \ + /* A * (B * s) */ \ + return scale_multiply(src.rhs().rhs().functor().m_other, \ + src.lhs(), src.rhs().lhs()); \ + } else { \ + return multiply(src.lhs(), src.rhs()); \ + } \ + }(); \ + if constexpr (is_add) { \ + dst += product; \ + } else if constexpr (is_sub) { \ + dst -= product; \ + } else { \ + dst = product; \ + } \ + } \ + } \ + }; #define STAN_EIGEN_AD_PRODUCT_ASSIGNMENT(TPL, SCALAR) \ STAN_EIGEN_AD_PRODUCT_ASSIGNMENT_OP(TPL, SCALAR, assign_op) \ diff --git a/stan/math/rev/core/Eigen_NumTraits.hpp b/stan/math/rev/core/Eigen_NumTraits.hpp index e3582600e6c..6ca7a84c293 100644 --- a/stan/math/rev/core/Eigen_NumTraits.hpp +++ b/stan/math/rev/core/Eigen_NumTraits.hpp @@ -302,6 +302,29 @@ struct ternary_evaluator< }; #endif +/** + * Coefficient-wise products and products with a scalar of `var` with `var`, + * `double` or `int` are computed by `elt_multiply`, `multiply` or + * `scale_multiply` (see stan/math/prim/meta/eigen_product_evaluator_hooks.hpp). + */ +#define STAN_REV_NO_TPARAMS() +STAN_EIGEN_AD_CWISE_PRODUCT_EVALUATOR(STAN_REV_NO_TPARAMS, stan::math::var, + stan::math::var) +STAN_EIGEN_AD_CWISE_PRODUCT_EVALUATOR(STAN_REV_NO_TPARAMS, stan::math::var, + double) +STAN_EIGEN_AD_CWISE_PRODUCT_EVALUATOR(STAN_REV_NO_TPARAMS, double, + stan::math::var) +STAN_EIGEN_AD_CWISE_PRODUCT_EVALUATOR(STAN_REV_NO_TPARAMS, stan::math::var, int) +STAN_EIGEN_AD_CWISE_PRODUCT_EVALUATOR(STAN_REV_NO_TPARAMS, int, stan::math::var) + +/** + * Products assigned to (`=`, `+=`, `-=`) a destination with `var` scalars are + * computed by `multiply`, `scale_multiply`, `multiply_add` or `subtract` (see + * stan/math/prim/meta/eigen_product_evaluator_hooks.hpp). + */ +STAN_EIGEN_AD_PRODUCT_ASSIGNMENT(STAN_REV_NO_TPARAMS, stan::math::var) +#undef STAN_REV_NO_TPARAMS + /** * Enable linear access of inputs when using read_vi_val_adj. */ @@ -504,6 +527,5 @@ struct general_matrix_matrix_product #endif diff --git a/stan/math/rev/core/eigen_product_hooks.hpp b/stan/math/rev/core/eigen_product_hooks.hpp deleted file mode 100644 index 74a7127021d..00000000000 --- a/stan/math/rev/core/eigen_product_hooks.hpp +++ /dev/null @@ -1,99 +0,0 @@ -#ifndef STAN_MATH_REV_CORE_EIGEN_PRODUCT_HOOKS_HPP -#define STAN_MATH_REV_CORE_EIGEN_PRODUCT_HOOKS_HPP - -#include -#include -#include -#include -#include -#include -#include - -/** - * @file - * Routes Eigen products whose operands have `var` or `std::complex` - * scalars (mixed with arithmetic scalars) to the reverse-mode kernels in - * stan/math/rev/core/eigen_product_kernel.hpp. - * - * - matrix products `A * B` (every Eigen product tag, `lazyProduct`, - * `noalias()`, `=`, `+=` and `-=`) through the `product_evaluator` and - * `Assignment` hooks of - * stan/math/prim/meta/eigen_product_evaluator_hooks.hpp; - * - coefficient-wise products `a.array() * b.array()`, `cwiseProduct`, - * and products with a scalar `s * A`, `A * s` through `binary_evaluator`. - * - * Scalar factors are folded into the product they scale, so `s * (A * B)`, - * `(A * B) * s` and `(s * A) * B` each put one callback on the stack. - * - * The product is evaluated in full when its evaluator is created, even when - * only some coefficients are read afterwards (e.g. `(a * b).head(2)`). An - * unread coefficient's adjoint is zero, but the callback still multiplies it - * by the operand values, so a non-finite operand value can turn it into NaN. - * In particular, a product inside a branch of `.select()` is evaluated for - * every coefficient, unlike other expressions there. - */ - -namespace stan { -namespace math { -namespace internal { - -template -inline constexpr bool is_rev_product_v - = is_rev_product_operand_scalar< - typename Eigen::internal::traits::Scalar>::value&& - is_rev_product_operand_scalar< - typename Eigen::internal::traits::Scalar>::value - && (is_rev_product_scalar< - typename Eigen::internal::traits::Scalar>::value - || is_rev_product_scalar< - typename Eigen::internal::traits::Scalar>::value); - -/** - * Reverse-mode matrix product. - */ -template -struct eigen_matrix_product_impl>> { - using result_t = typename Eigen::Product::PlainObject; - static result_t run(const Lhs& lhs, const Rhs& rhs) { - return rev_matrix_product(lhs, rhs); - } -}; - -/** - * Reverse-mode coefficient-wise product, or product with a scalar. - */ -template -struct eigen_cwise_product_impl< - Eigen::CwiseBinaryOp, Lhs, Rhs>, - std::enable_if_t>> { - using xpr_t = Eigen::CwiseBinaryOp, - Lhs, Rhs>; - using result_t = typename xpr_t::PlainObject; - static result_t run(const xpr_t& xpr) { - return rev_cwise_product(xpr); - } -}; - -} // namespace internal -} // namespace math -} // namespace stan - -namespace Eigen { -namespace internal { - -#define STAN_REV_PRODUCT_NO_TPARAMS() -#define STAN_REV_PRODUCT_S_TPARAMS() typename S_, -STAN_EIGEN_AD_SCALAR_PATTERNS(STAN_EIGEN_AD_CWISE_PRODUCT_EVALUATOR, - STAN_REV_PRODUCT_NO_TPARAMS, - STAN_REV_PRODUCT_S_TPARAMS, stan::math::var) -STAN_EIGEN_AD_PRODUCT_ASSIGNMENT(STAN_REV_PRODUCT_NO_TPARAMS, stan::math::var) -STAN_EIGEN_AD_PRODUCT_ASSIGNMENT(STAN_REV_PRODUCT_NO_TPARAMS, - std::complex) -#undef STAN_REV_PRODUCT_NO_TPARAMS -#undef STAN_REV_PRODUCT_S_TPARAMS - -} // namespace internal -} // namespace Eigen - -#endif diff --git a/stan/math/rev/core/eigen_product_kernel.hpp b/stan/math/rev/core/eigen_product_kernel.hpp deleted file mode 100644 index a6b43a1cb87..00000000000 --- a/stan/math/rev/core/eigen_product_kernel.hpp +++ /dev/null @@ -1,440 +0,0 @@ -#ifndef STAN_MATH_REV_CORE_EIGEN_PRODUCT_KERNEL_HPP -#define STAN_MATH_REV_CORE_EIGEN_PRODUCT_KERNEL_HPP - -#include -#include -#include -#include -#include -#include -#include -#include -#include - -/** - * @file - * Reverse-mode kernels for Eigen products whose operands have `var` or - * `std::complex` scalars. Each product puts one callback on the - * reverse-mode stack plus one value-only (nochain) vari per result - * coefficient, instead of one vari per scalar operation. - * - * All arithmetic inside the kernels is on `double` or `std::complex` - * values. The kernels must never form an Eigen product of `var` scalars - * themselves, since that would dispatch back into these kernels. - */ - -namespace stan { -namespace math { -namespace internal { - -/** - * True for `var` and `std::complex`. - */ -template -struct is_rev_product_scalar : std::false_type {}; -template <> -struct is_rev_product_scalar : std::true_type {}; -template <> -struct is_rev_product_scalar> : std::true_type {}; - -/** - * True for scalars a reverse-mode product may mix with `var`. - */ -template -struct is_rev_product_operand_scalar - : bool_constant::value - || is_rev_product_scalar::value> {}; -template -struct is_rev_product_operand_scalar> - : bool_constant::value - || std::is_same::value> {}; - -/** - * The value type used for the arithmetic of a product whose result has - * scalar type `T`: `double`, or `std::complex` for complex results. - */ -template -struct rev_product_value { - using type = double; -}; -template -struct rev_product_value> { - using type = std::complex; -}; -template -using rev_product_value_t = typename rev_product_value::type; - -template * = nullptr> -inline double rev_product_value_of(const T& x) { - return x; -} -template * = nullptr> -inline std::complex rev_product_value_of(const std::complex& z) { - return {static_cast(z.real()), static_cast(z.imag())}; -} -inline double rev_product_value_of(const var& x) { return x.val(); } -inline std::complex rev_product_value_of(const std::complex& z) { - return {z.real().val(), z.imag().val()}; -} - -inline double rev_product_conj(double x) { return x; } -inline std::complex rev_product_conj(const std::complex& z) { - return std::conj(z); -} - -/** - * Add the adjoint `g` of a product result to operand `x`. For a real `x` and - * a complex `g`, only the real part contributes. - */ -inline void rev_product_add_adjoint(const var& x, double g) { - x.vi_->adj_ += g; -} -inline void rev_product_add_adjoint(const var& x, - const std::complex& g) { - x.vi_->adj_ += g.real(); -} -inline void rev_product_add_adjoint(const std::complex& z, - const std::complex& g) { - z.real().vi_->adj_ += g.real(); - z.imag().vi_->adj_ += g.imag(); -} - -/** - * A new value-only vari holding `x`. Results of a product kernel get their - * adjoints from the kernel's callback, so they need no `chain()`. - */ -inline var rev_product_result(double x) { return var(new vari(x, false)); } -inline std::complex rev_product_result(const std::complex& z) { - return {var(new vari(z.real(), false)), var(new vari(z.imag(), false))}; -} - -/** - * Operand of a reverse-mode product kernel: its values, and for autodiff - * scalars its vars, copied to the arena once. - * - * @tparam Scalar scalar type of the operand expression - */ -template ::value> -struct rev_product_operand { - static constexpr bool is_autodiff = false; - using value_t = rev_product_value_t; - arena_matrix> val_; - template - explicit rev_product_operand(const X& x) : val_(x.template cast()) {} - const auto& val() const { return val_; } - template - void add_adjoint(const G&) const {} -}; - -template -struct rev_product_operand { - static constexpr bool is_autodiff = true; - using value_t = rev_product_value_t; - arena_matrix> vars_; - arena_matrix> val_; - template - explicit rev_product_operand(const X& x) - : vars_(x), val_(vars_.unaryExpr([](const Scalar& s) { - return rev_product_value_of(s); - })) {} - const auto& val() const { return val_; } - template - void add_adjoint(const G& g) const { - using g_scalar = typename Eigen::internal::traits::Scalar; - if constexpr (std::is_same::value - && std::is_same::value) { - // a shallow copy of the arena map, whose adjoint view is writable - arena_matrix> vars - = vars_; - vars.adj() += g; - } else { - const auto& g_ref = g.eval(); - for (Eigen::Index j = 0; j < vars_.cols(); ++j) { - for (Eigen::Index i = 0; i < vars_.rows(); ++i) { - rev_product_add_adjoint(vars_.coeff(i, j), g_ref.coeff(i, j)); - } - } - } - } -}; - -/** - * Scalar factors folded into a product kernel: the product of the - * arithmetic factors, and the autodiff factors. - * - * @tparam ValueT `double` or `std::complex` - */ -template -struct rev_product_scales { - // three peels (expression, left and right operand) of up to - // max_peel_depth factors each - static constexpr int max_autodiff = 3 * max_peel_depth; - struct autodiff_scale { - vari* re_; - vari* im_; // nullptr for a real scale - ValueT value() const { - if constexpr (std::is_same::value) { - return re_->val_; - } else { - return {re_->val_, im_ == nullptr ? 0.0 : im_->val_}; - } - } - }; - ValueT constant_{1.0}; - // allocated on the arena at the first autodiff factor, so a kernel's - // callback captures a pointer instead of storage for every factor - autodiff_scale* autodiff_{nullptr}; - int n_autodiff_{0}; - - template * = nullptr> - void push(const T& x) { - constant_ *= static_cast(x); - } - template * = nullptr> - void push(const std::complex& x) { - constant_ *= std::complex(x.real(), x.imag()); - } - void push(const var& x) { push_autodiff({x.vi_, nullptr}); } - void push(const std::complex& x) { - push_autodiff({x.real().vi_, x.imag().vi_}); - } - - /** Product of every factor. */ - ValueT value() const { - ValueT v = constant_; - for (int i = 0; i < n_autodiff_; ++i) { - v *= autodiff_[i].value(); - } - return v; - } - - /** - * Given `t = sum(G .* conj(M))` for the unscaled product `M` and result - * adjoints `G`, add each autodiff factor's adjoint. - */ - void add_adjoints(const ValueT& t) const { - for (int i = 0; i < n_autodiff_; ++i) { - ValueT others = constant_; - for (int k = 0; k < n_autodiff_; ++k) { - if (k != i) { - others *= autodiff_[k].value(); - } - } - const ValueT g = rev_product_conj(others) * t; - if constexpr (std::is_same::value) { - autodiff_[i].re_->adj_ += g; - } else { - autodiff_[i].re_->adj_ += g.real(); - if (autodiff_[i].im_ != nullptr) { - autodiff_[i].im_->adj_ += g.imag(); - } - } - } - } - - private: - void push_autodiff(const autodiff_scale& s) { - if (autodiff_ == nullptr) { - autodiff_ - = ChainableStack::instance_->memalloc_.alloc_array( - max_autodiff); - } - autodiff_[n_autodiff_++] = s; - } -}; - -template -using rev_result_value_t = rev_product_value_t; - -/** - * Build the plain result of a product kernel from its values, with one - * nochain vari per coefficient, and keep an arena copy for the callback. - */ -template -inline Result rev_product_make_result( - const Values& values, - arena_matrix>& res_arena) { - res_arena = arena_matrix< - Eigen::Matrix>( - values.rows(), values.cols()); - Result res(values.rows(), values.cols()); - for (Eigen::Index j = 0; j < values.cols(); ++j) { - for (Eigen::Index i = 0; i < values.rows(); ++i) { - res_arena.coeffRef(i, j) = rev_product_result(values.coeff(i, j)); - res.coeffRef(i, j) = res_arena.coeff(i, j); - } - } - return res; -} - -/** - * Adjoints of a kernel result as a lazy `double` or `std::complex` - * expression. - */ -template -inline auto rev_product_result_adjoints( - const arena_matrix< - Eigen::Matrix>& res) { - if constexpr (std::is_same::value) { - return res.adj(); - } else { - return res.unaryExpr([](const ResScalar& z) { - return std::complex(z.real().adj(), z.imag().adj()); - }); - } -} - -template -using rev_operand_t - = rev_product_operand::Scalar>; - -/** - * `sigma * X * Y` for matrix operands `X` and `Y`, with one callback. - * - * @tparam Result plain Eigen type of the product - */ -template -inline Result rev_matmul(const Scales& scales, const XOp& x, const YOp& y) { - using value_t = rev_result_value_t; - using res_scalar = typename Result::Scalar; - const value_t sigma = scales.value(); - arena_matrix> - res_arena; - arena_matrix> unscaled; - Result res; - if (scales.n_autodiff_ > 0) { - unscaled = x.val() * y.val(); - res = rev_product_make_result((sigma * unscaled).eval(), res_arena); - } else { - res = rev_product_make_result((sigma * (x.val() * y.val())).eval(), - res_arena); - } - reverse_pass_callback([res_arena, x, y, scales, sigma, unscaled]() { - // evaluated once: both operand gradients are products with it - const auto g = rev_product_result_adjoints(res_arena).eval(); - const value_t sigma_conj = rev_product_conj(sigma); - if constexpr (XOp::is_autodiff) { - x.add_adjoint(sigma_conj * (g * y.val().adjoint())); - } - if constexpr (YOp::is_autodiff) { - y.add_adjoint(sigma_conj * (x.val().adjoint() * g)); - } - if (scales.n_autodiff_ > 0) { - scales.add_adjoints((g.array() * unscaled.array().conjugate()).sum()); - } - }); - return res; -} - -/** - * `sigma * (X .* Y)` for operands of equal size, with one callback. - */ -template -inline Result rev_cwise(const Scales& scales, const XOp& x, const YOp& y) { - using value_t = rev_result_value_t; - using res_scalar = typename Result::Scalar; - const value_t sigma = scales.value(); - arena_matrix> - res_arena; - Result res = rev_product_make_result( - (sigma * (x.val().array() * y.val().array())).matrix().eval(), res_arena); - reverse_pass_callback([res_arena, x, y, scales, sigma]() { - const auto g = rev_product_result_adjoints(res_arena); - const value_t sigma_conj = rev_product_conj(sigma); - if constexpr (XOp::is_autodiff) { - x.add_adjoint( - (sigma_conj * (g.array() * y.val().array().conjugate())).matrix()); - } - if constexpr (YOp::is_autodiff) { - y.add_adjoint( - (sigma_conj * (g.array() * x.val().array().conjugate())).matrix()); - } - if (scales.n_autodiff_ > 0) { - scales.add_adjoints( - (g.array() * (x.val().array() * y.val().array()).conjugate()).sum()); - } - }); - return res; -} - -/** - * `sigma * X` for a matrix operand `X`, with one callback. - */ -template -inline Result rev_scale(const Scales& scales, const XOp& x) { - using value_t = rev_result_value_t; - using res_scalar = typename Result::Scalar; - const value_t sigma = scales.value(); - arena_matrix> - res_arena; - Result res - = rev_product_make_result((sigma * x.val()).eval(), res_arena); - reverse_pass_callback([res_arena, x, scales, sigma]() { - const auto g = rev_product_result_adjoints(res_arena); - if constexpr (XOp::is_autodiff) { - x.add_adjoint(rev_product_conj(sigma) * g); - } - if (scales.n_autodiff_ > 0) { - scales.add_adjoints((g.array() * x.val().array().conjugate()).sum()); - } - }); - return res; -} - -/** - * Matrix product `lhs * rhs` of Eigen expressions. Scalar factors of either - * operand are folded into the product's callback. - * - * @tparam Result plain Eigen type of the product - */ -template -inline Result rev_matrix_product(const Lhs& lhs, const Rhs& rhs) { - rev_product_scales> scales; - const auto& lhs_core = peel_scales(scales, lhs); - const auto& rhs_core = peel_scales(scales, rhs); - using lhs_core_t = std::decay_t; - using rhs_core_t = std::decay_t; - return rev_matmul(scales, rev_operand_t(lhs_core), - rev_operand_t(rhs_core)); -} - -/** - * Coefficient-wise product expression `xpr` (including a matrix times a - * scalar). A scaled matrix product `(A * B) * s` becomes one matrix-product - * callback. - * - * @tparam Result plain Eigen type of the expression - */ -template -inline Result rev_cwise_product(const Xpr& xpr) { - using value_t = rev_result_value_t; - rev_product_scales scales; - const auto& core = peel_scales(scales, xpr); - using core_t = std::decay_t; - if constexpr (std::is_same::value) { - // a coefficient-wise product of two non-constant operands - const auto& lhs_core = peel_scales(scales, xpr.lhs()); - const auto& rhs_core = peel_scales(scales, xpr.rhs()); - using lhs_core_t = std::decay_t; - using rhs_core_t = std::decay_t; - return rev_cwise(scales, rev_operand_t(lhs_core), - rev_operand_t(rhs_core)); - } else if constexpr (is_dense_product::value) { - const auto& lhs_core = peel_scales(scales, core.lhs()); - const auto& rhs_core = peel_scales(scales, core.rhs()); - using lhs_core_t = std::decay_t; - using rhs_core_t = std::decay_t; - return rev_matmul(scales, rev_operand_t(lhs_core), - rev_operand_t(rhs_core)); - } else { - return rev_scale(scales, rev_operand_t(core)); - } -} - -} // namespace internal -} // namespace math -} // namespace stan - -#endif diff --git a/stan/math/rev/core/var.hpp b/stan/math/rev/core/var.hpp index 8306fc8a2a0..7186990ca98 100644 --- a/stan/math/rev/core/var.hpp +++ b/stan/math/rev/core/var.hpp @@ -8,7 +8,6 @@ #include #include #include -#include #include #include #include @@ -1187,16 +1186,6 @@ class var_value> { // For backwards compatibility the default value is double using var = var_value; -namespace internal { -/** - * Eigen products of `var` scalars are evaluated by Stan (see - * stan/math/prim/meta/eigen_product_evaluator_hooks.hpp). Declared with - * `var` itself so every translation unit that uses `var` agrees. - */ -template <> -struct is_stan_ad_product_scalar : std::true_type {}; -} // namespace internal - } // namespace math /** diff --git a/stan/math/rev/fun.hpp b/stan/math/rev/fun.hpp index d6f02c3f13f..8f27e23b1dd 100644 --- a/stan/math/rev/fun.hpp +++ b/stan/math/rev/fun.hpp @@ -126,6 +126,7 @@ #include #include #include +#include #include #include #include @@ -149,6 +150,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..8a4d310d21e --- /dev/null +++ b/stan/math/rev/fun/multiply_add.hpp @@ -0,0 +1,69 @@ +#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(const Mat1& A, const Mat2& B, const 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 = A; + arena_t arena_B = B; + arena_t arena_C = C; + auto arena_A_val = to_arena(value_of(arena_A)); + auto arena_B_val = to_arena(value_of(arena_B)); + using return_t = return_var_matrix_t; + arena_t res = arena_A_val * arena_B_val + value_of(arena_C); + reverse_pass_callback( + [arena_A, arena_B, arena_C, arena_A_val, arena_B_val, res]() mutable { + const auto res_adj = res.adj().eval(); + if constexpr (is_autodiff_v) { + arena_A.adj() += res_adj * arena_B_val.transpose(); + } + if constexpr (is_autodiff_v) { + arena_B.adj() += arena_A_val.transpose() * res_adj; + } + if constexpr (is_autodiff_v) { + arena_C.adj() += res_adj; + } + }); + return return_t(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..a7f45add5fa --- /dev/null +++ b/stan/math/rev/fun/scale_multiply.hpp @@ -0,0 +1,66 @@ +#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(const Scal& s, const Mat1& A, const Mat2& B) { + check_multiplicable("scale_multiply", "A", A, "B", B); + arena_t arena_A = A; + arena_t arena_B = B; + auto arena_A_val = to_arena(value_of(arena_A)); + auto arena_B_val = to_arena(value_of(arena_B)); + const double s_val = value_of(s); + using return_t + = return_var_matrix_t; + arena_t res = s_val * (arena_A_val * arena_B_val); + reverse_pass_callback( + [s, arena_A, arena_B, arena_A_val, arena_B_val, s_val, res]() mutable { + const auto res_adj = res.adj().eval(); + if constexpr (is_autodiff_v) { + s.adj() += (arena_A_val.transpose() * res_adj) + .cwiseProduct(arena_B_val) + .sum(); + } + if constexpr (is_autodiff_v) { + arena_A.adj() += s_val * res_adj * arena_B_val.transpose(); + } + if constexpr (is_autodiff_v) { + arena_B.adj() += s_val * arena_A_val.transpose() * res_adj; + } + }); + return return_t(res); +} + +} // namespace math +} // namespace stan + +#endif diff --git a/test/unit/math/fwd/core/eigen_product_hook_test.cpp b/test/unit/math/fwd/core/eigen_product_hook_test.cpp deleted file mode 100644 index a74d8b804a7..00000000000 --- a/test/unit/math/fwd/core/eigen_product_hook_test.cpp +++ /dev/null @@ -1,82 +0,0 @@ -#include -#include -#include -#include - -// Eigen products of fvar matrices go through Stan's product hooks: -// matrix products are split into double products of values and tangents, -// coefficient-wise products stay lazy. - -namespace { -using stan::math::fvar; -using fvar_d = fvar; -using stan::math::internal::is_stan_ad_product_hook_v; - -template -Eigen::Matrix fvar_mat(int rows, int cols, double offset) { - Eigen::Matrix x(rows, cols); - for (int j = 0; j < cols; ++j) { - for (int i = 0; i < rows; ++i) { - x(i, j) = fvar_d(0.1 * (i + 1) - 0.2 * j + offset, 0.3 * j - 0.1 * i); - } - } - return x; -} - -template -void expect_fvar_product(const A& a, const B& b) { - Eigen::Matrix c = a * b; - ASSERT_EQ(c.rows(), a.rows()); - ASSERT_EQ(c.cols(), b.cols()); - for (int j = 0; j < c.cols(); ++j) { - for (int i = 0; i < c.rows(); ++i) { - fvar_d expected = 0; - for (int k = 0; k < a.cols(); ++k) { - expected += a(i, k) * b(k, j); - } - EXPECT_NEAR(c(i, j).val_, expected.val_, 1e-12); - EXPECT_NEAR(c(i, j).d_, expected.d_, 1e-12); - } - } -} -} // namespace - -TEST(mathFwdEigenProductHook, selected) { - using m = Eigen::Matrix; - using md = Eigen::MatrixXd; - EXPECT_TRUE((is_stan_ad_product_hook_v() * std::declval())>>)); - EXPECT_TRUE((is_stan_ad_product_hook_v() * std::declval())>>)); - EXPECT_TRUE((is_stan_ad_product_hook_v< - std::decay_t().array() - * std::declval().array())>>)); - EXPECT_FALSE((is_stan_ad_product_hook_v() * std::declval())>>)); -} - -TEST(mathFwdEigenProductHook, matmulValuesAndTangents) { - expect_fvar_product(fvar_mat(3, 4, 0.1), fvar_mat(4, 2, -0.2)); - const Eigen::MatrixXd d = Eigen::MatrixXd::Constant(4, 3, 0.5); - expect_fvar_product(fvar_mat(2, 4, 0.1), d); - expect_fvar_product(d.transpose(), fvar_mat(4, 1, 0.3)); - expect_fvar_product(fvar_mat<2, 2>(2, 2, 0.1), fvar_mat<2, 2>(2, 2, 0.4)); -} - -TEST(mathFwdEigenProductHook, cwiseIsLazy) { - auto a = fvar_mat(3, 2, 0.1); - auto b = fvar_mat(3, 2, -0.4); - auto xpr = a.array() * b.array(); - Eigen::internal::evaluator eval(xpr); - // lazy evaluation: the evaluator does not own a plain result - EXPECT_FALSE( - static_cast(Eigen::internal::evaluator::Flags - & Eigen::EvalBeforeNestingBit)); - for (int j = 0; j < 2; ++j) { - for (int i = 0; i < 3; ++i) { - const fvar_d expected = a(i, j) * b(i, j); - EXPECT_NEAR(eval.coeff(i, j).val_, expected.val_, 1e-12); - EXPECT_NEAR(eval.coeff(i, j).d_, expected.d_, 1e-12); - } - } -} diff --git a/test/unit/math/mix/core/eigen_product_hook_tape_test.cpp b/test/unit/math/mix/core/eigen_product_hook_tape_test.cpp deleted file mode 100644 index f9e00b77902..00000000000 --- a/test/unit/math/mix/core/eigen_product_hook_tape_test.cpp +++ /dev/null @@ -1,72 +0,0 @@ -#include -#include -#include -#include - -// Reverse-mode nodes recorded by Eigen products of fvar matrices, -// which Stan splits into products of their var values and tangents. - -namespace { -using stan::math::fvar; -using stan::math::var; -using fvar_v = fvar; - -Eigen::Matrix fvar_var_mat(int rows, int cols, double offset) { - Eigen::Matrix x(rows, cols); - for (int j = 0; j < cols; ++j) { - for (int i = 0; i < rows; ++i) { - x(i, j) = fvar_v(var(0.1 * (i + 1) - 0.2 * j + offset), var(1.0)); - } - } - return x; -} -} // namespace - -TEST(mathMixEigenProductHook, fvarVarProductsAreHooked) { - using stan::math::internal::is_stan_ad_product_hook_v; - using m = Eigen::Matrix; - using md = Eigen::MatrixXd; - EXPECT_TRUE((is_stan_ad_product_hook_v() * std::declval())>>)); - EXPECT_TRUE((is_stan_ad_product_hook_v() * std::declval())>>)); - EXPECT_TRUE((is_stan_ad_product_hook_v< - std::decay_t().array() - * std::declval().array())>>)); -} - -TEST(mathMixEigenProductHook, fvarVarMatmulTape) { - stan::math::nested_rev_autodiff nested; - const int n = 6, k = 5, m = 4; - auto a = fvar_var_mat(n, k, 0.1); - auto b = fvar_var_mat(k, m, -0.3); - const stan::test::tape_delta delta; - Eigen::Matrix c = a * b; - // one callback for the values and one for the concatenated tangent product - EXPECT_EQ(delta.chain(), 2); - EXPECT_EQ(delta.nochain(), 2 * n * m); -} - -TEST(mathMixEigenProductHook, fvarVarDoubleMatmulTape) { - stan::math::nested_rev_autodiff nested; - const int n = 6, k = 5; - const Eigen::MatrixXd x = Eigen::MatrixXd::Constant(n, k, 0.25); - auto beta = fvar_var_mat(k, 1, 0.2); - const stan::test::tape_delta delta; - Eigen::Matrix xb = x * beta; - EXPECT_EQ(delta.chain(), 2); - EXPECT_EQ(delta.nochain(), 2 * n); -} - -TEST(mathMixEigenProductHook, fvarVarCwiseTape) { - stan::math::nested_rev_autodiff nested; - const int n = 7; - auto a = fvar_var_mat(n, 1, 0.1); - auto b = fvar_var_mat(n, 1, -0.3); - const stan::test::tape_delta delta; - Eigen::Matrix c = (a.array() * b.array()).matrix(); - // values: one callback; tangent: two callbacks plus one vari per - // coefficient for their sum - EXPECT_EQ(delta.chain(), 3 + n); - EXPECT_EQ(delta.nochain(), 3 * n); -} 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_complex_test.cpp b/test/unit/math/mix/fun/multiply_operator_complex_test.cpp deleted file mode 100644 index 31d2ef05658..00000000000 --- a/test/unit/math/mix/fun/multiply_operator_complex_test.cpp +++ /dev/null @@ -1,49 +0,0 @@ -#include -#include - -// Eigen products of complex autodiff matrices are evaluated by Stan's -// product hooks. - -using stan::test::multiply_operator::expect_hooked; - -namespace { -Eigen::MatrixXcd complex_matrix() { - Eigen::MatrixXcd x(2, 2); - x << std::complex(1, 0.5), 2, std::complex(3, -1), 4; - return x; -} -} // namespace - -TEST(mathMixMultiplyOperator, complexMatrixProduct) { - auto xpr = [](const auto& x, const auto& y) { return x * y; }; - auto f = [](const auto& x, const auto& y) { return (x * y).eval(); }; - Eigen::MatrixXd d_mat(2, 2); - d_mat << 1, 2, 3, 4; - Eigen::MatrixXcd cd_mat = complex_matrix(); - Eigen::VectorXcd cd_vec(2); - cd_vec << std::complex(0.5, 1), 2; - expect_hooked(xpr, cd_mat, cd_mat); - expect_hooked(xpr, cd_mat, cd_vec); - stan::test::expect_ad(f, d_mat, cd_mat); - stan::test::expect_ad(f, cd_mat, d_mat); - stan::test::expect_ad(f, cd_mat, cd_mat); - stan::test::expect_ad(f, cd_mat, cd_vec); -} - -TEST(mathMixMultiplyOperator, complexCwiseAndScalar) { - auto cwise = [](const auto& x, const auto& y) { return x.cwiseProduct(y); }; - auto scale = [](const auto& s, const auto& x) { return s * x; }; - auto f_cwise - = [](const auto& x, const auto& y) { return x.cwiseProduct(y).eval(); }; - auto f_scale = [](const auto& s, const auto& x) { return (s * x).eval(); }; - Eigen::MatrixXd d_mat(2, 2); - d_mat << 1, 2, 3, 4; - Eigen::MatrixXcd cd_mat = complex_matrix(); - std::complex cd_scalar(0.5, -0.25); - expect_hooked(cwise, cd_mat, cd_mat); - expect_hooked(scale, cd_scalar, cd_mat); - stan::test::expect_ad(f_cwise, cd_mat, cd_mat); - stan::test::expect_ad(f_cwise, d_mat, cd_mat); - stan::test::expect_ad(f_scale, cd_scalar, d_mat); - stan::test::expect_ad(f_scale, cd_scalar, cd_mat); -} diff --git a/test/unit/math/mix/fun/multiply_operator_cwise_test.cpp b/test/unit/math/mix/fun/multiply_operator_cwise_test.cpp index faece5b1c76..45e850871a4 100644 --- a/test/unit/math/mix/fun/multiply_operator_cwise_test.cpp +++ b/test/unit/math/mix/fun/multiply_operator_cwise_test.cpp @@ -1,30 +1,26 @@ #include #include -// Coefficient-wise products and products with a scalar of autodiff -// matrices and arrays are evaluated by Stan's product hooks. +// 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::expect_hooked; using stan::test::multiply_operator::test_matrix; TEST(mathMixMultiplyOperator, arrayTimesArray) { - auto xpr = [](const auto& a, const auto& b) { return a.array() * b.array(); }; 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); - expect_hooked(xpr, a, b); stan::test::expect_ad(f, a, b); } TEST(mathMixMultiplyOperator, cwiseProduct) { - auto xpr = [](const auto& a, const auto& b) { return a.cwiseProduct(b); }; 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); - expect_hooked(xpr, a, b); stan::test::expect_ad(f, a, b); } @@ -42,36 +38,19 @@ TEST(mathMixMultiplyOperator, cwiseNestedAndPartialRead) { } TEST(mathMixMultiplyOperator, scalarTimesMatrix) { - auto left = [](const auto& s, const auto& a) { return s * a; }; - auto right = [](const auto& s, const auto& a) { return a * s; }; - auto array = [](const auto& s, const auto& a) { return a.array() * s; }; + 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); - expect_hooked(left, s, a); - expect_hooked(right, s, a); - expect_hooked(array, s, a); - auto eval = [](const auto& f) { - return [f](const auto& s, const auto& a) { - return Eigen::Matrix, -1, - -1>(f(s, a)); - }; - }; - stan::test::expect_ad(eval(left), s, a); - stan::test::expect_ad(eval(right), s, a); - stan::test::expect_ad(eval(array), s, a); + stan::test::expect_ad(left, s, a); + stan::test::expect_ad(right, s, a); + stan::test::expect_ad(array, s, a); } TEST(mathMixMultiplyOperator, matrixTimesInt) { - Eigen::MatrixXd a = test_matrix(2, 3, 0.2); - // only the matrix can be autodiff here, so check the left promotion - using stan::math::fvar; - using stan::math::var; - EXPECT_TRUE( - (stan::math::internal::is_stan_ad_product_hook_v&>() * 3)>>)); - EXPECT_TRUE( - (stan::math::internal::is_stan_ad_product_hook_v, -1, -1>&>() * 3)>>)); auto f = [](const auto& a) { return (a * 3).eval(); }; - stan::test::expect_ad(f, a); + 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 index f159db1c5e5..2314f91200d 100644 --- a/test/unit/math/mix/fun/multiply_operator_matmul_test.cpp +++ b/test/unit/math/mix/fun/multiply_operator_matmul_test.cpp @@ -1,58 +1,48 @@ #include #include -// Eigen `operator*` matrix products of autodiff matrices are evaluated by -// Stan's product hooks. Each test checks that the hook is selected for every -// autodiff scalar and that values and derivatives are correct. +// 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::expect_assignment_hooked; -using stan::test::multiply_operator::expect_hooked; using stan::test::multiply_operator::test_matrix; namespace { -auto product = [](const auto& a, const auto& b) { return a * b; }; +// 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); - expect_hooked(product, a, b); - stan::test::expect_ad(product_eval, a, b); + stan::test::expect_ad(product_eval, a, b); Eigen::MatrixXd empty(0, 4); - stan::test::expect_ad(product_eval, empty, b); + 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); - expect_hooked(product, a, v); - stan::test::expect_ad(product_eval, a, v); + 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); - expect_hooked(product, rv, a); - stan::test::expect_ad(product_eval, rv, a); + 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); - expect_hooked(product, rv, v); - expect_hooked(product, v, rv); - stan::test::expect_ad(product_eval, rv, v); - stan::test::expect_ad(product_eval, v, rv); + stan::test::expect_ad(product_eval, rv, v); + stan::test::expect_ad(product_eval, v, rv); } TEST(mathMixMultiplyOperator, fixedSize) { - // 2x2 uses Eigen's lazy coefficient-based product, 4x4 the - // coefficient-based one - expect_hooked(product, Eigen::Matrix2d(test_matrix(2, 2, 0.1)), - Eigen::Matrix2d(test_matrix(2, 2, -0.3))); - expect_hooked(product, Eigen::Matrix4d(test_matrix(4, 4, 0.1)), - Eigen::Matrix4d(test_matrix(4, 4, -0.3))); + // 2x2 uses Eigen's lazy coefficient-based product auto f_fixed = [](const auto& a, const auto& b) { using a_scalar = typename std::decay_t::Scalar; using b_scalar = typename std::decay_t::Scalar; @@ -62,38 +52,32 @@ TEST(mathMixMultiplyOperator, fixedSize) { const Eigen::Matrix c = a2 * b2; return Eigen::Matrix(c); }; - stan::test::expect_ad(f_fixed, test_matrix(2, 2, 0.1), - test_matrix(2, 2, -0.3)); + stan::test::expect_ad(f_fixed, test_matrix(2, 2, 0.1), + test_matrix(2, 2, -0.3)); } TEST(mathMixMultiplyOperator, lazyProduct) { - auto lazy = [](const auto& a, const auto& b) { return a.lazyProduct(b); }; 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); - expect_hooked(lazy, a, b); - stan::test::expect_ad(lazy_eval, a, b); + 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); + return (a.transpose() * b.block(0, 0, 3, 2)).eval(); }; - auto f_t_eval - = [&f_t](const auto& a, const auto& b) { return f_t(a, b).eval(); }; Eigen::MatrixXd a = test_matrix(3, 4, 0.1); Eigen::MatrixXd b = test_matrix(3, 3, -0.2); - expect_hooked(f_t, a, b); - stan::test::expect_ad(f_t_eval, a, b); + 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)); + stan::test::expect_ad(f_chain, a, test_matrix(4, 2, 0.4)); } TEST(mathMixMultiplyOperator, assignment) { - expect_assignment_hooked(); auto f_add_sub = [](const auto& a, const auto& b) { auto c = (a * b).eval(); c += a * b; @@ -104,12 +88,13 @@ TEST(mathMixMultiplyOperator, assignment) { 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); + stan::test::expect_ad(f_add_sub, a, b); + stan::test::expect_ad(f_noalias, a, b); } TEST(mathMixMultiplyOperator, triangularTimesVector) { diff --git a/test/unit/math/mix/fun/multiply_operator_scaled_test.cpp b/test/unit/math/mix/fun/multiply_operator_scaled_test.cpp index 5937bb9cd48..2a792024bc1 100644 --- a/test/unit/math/mix/fun/multiply_operator_scaled_test.cpp +++ b/test/unit/math/mix/fun/multiply_operator_scaled_test.cpp @@ -1,30 +1,12 @@ #include #include -// Scaled matrix products `s * (A * B)`, `(A * B) * s` and `(s * A) * B` are -// evaluated by Stan's product hooks with the scalar folded into the product. +// 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::expect_hooked; using stan::test::multiply_operator::test_matrix; -TEST(mathMixMultiplyOperator, scaledProductsHooked) { - const Eigen::MatrixXd b = test_matrix(2, 2, -0.3); - const Eigen::MatrixXd a = test_matrix(3, 2, 0.1); - double s = 0.7; - // autodiff scale and left operand, double right operand - expect_hooked([&b](const auto& s, const auto& a) { return s * (a * b); }, s, - a); - expect_hooked([&b](const auto& s, const auto& a) { return (a * b) * s; }, s, - a); - expect_hooked([&b](const auto& s, const auto& a) { return (s * a) * b; }, s, - a); - // autodiff scale and right operand, double left operand - expect_hooked([&a](const auto& s, const auto& b) { return s * (a * b); }, s, - b); - expect_hooked([&a](const auto& s, const auto& b) { return a * (b * s); }, s, - b); -} - TEST(mathMixMultiplyOperator, scaledProducts) { Eigen::MatrixXd a = test_matrix(3, 2, 0.1); Eigen::MatrixXd b = test_matrix(2, 2, -0.3); @@ -41,8 +23,11 @@ TEST(mathMixMultiplyOperator, scaledProducts) { 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(); }; - 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); + // 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 index 7b139907f6a..ec4a0883afc 100644 --- a/test/unit/math/mix/fun/multiply_operator_util.hpp +++ b/test/unit/math/mix/fun/multiply_operator_util.hpp @@ -2,111 +2,11 @@ #define TEST_UNIT_MATH_MIX_FUN_MULTIPLY_OPERATOR_UTIL_HPP #include -#include -#include -#include namespace stan { namespace test { namespace multiply_operator { -using stan::math::fvar; -using stan::math::var; - -/** - * The autodiff scalars `expect_ad` instantiates, for which Eigen products - * must be evaluated by Stan's product hooks. - */ -using ad_scalars = std::tuple, fvar, fvar>, - fvar>>; - -template -struct promote { - using type = stan::math::promote_scalar_t; -}; -template -struct promote { - using type = T; -}; -template -struct promote> { - using type = std::complex; -}; -template -struct promote, R, C, O, MR, MC>> { - using type = Eigen::Matrix, R, C, O, MR, MC>; -}; -template -using promote_t = typename promote::type; - -/** - * Expect Eigen to evaluate `xpr(a, b)` with a Stan product hook when either - * or both arguments are promoted to autodiff scalar `T`. `xpr` returns the - * unevaluated product expression; it is only used in `decltype`. - */ -template -void expect_hooked_for(const F& xpr, const A&, const B&) { - using a_t = promote_t; - using b_t = promote_t; - using both_t = std::decay_t(), std::declval()))>; - using left_t = std::decay_t(), std::declval()))>; - using right_t = std::decay_t(), std::declval()))>; - EXPECT_TRUE(stan::math::internal::is_stan_ad_product_hook_v) - << "both " << typeid(T).name(); - EXPECT_TRUE(stan::math::internal::is_stan_ad_product_hook_v) - << "left " << typeid(T).name(); - EXPECT_TRUE(stan::math::internal::is_stan_ad_product_hook_v) - << "right " << typeid(T).name(); -} - -/** - * `expect_hooked_for` for every autodiff scalar, and expect the - * all-`double` expression to stay with Eigen. - */ -template -void expect_hooked(const F& xpr, const A& a, const B& b) { - std::apply( - [&](auto... t) { (expect_hooked_for(xpr, a, b), ...); }, - ad_scalars{}); - using double_t = std::decay_t; - EXPECT_FALSE(stan::math::internal::is_stan_ad_product_hook_v); -} - -/** - * Expect Eigen's assignment of a product to a destination with autodiff - * scalar `T` (`=`, `+=`, `-=`, `noalias()`) to use Stan's hook, which - * `Eigen::internal::Assignment` otherwise bypasses. - */ -template -void expect_assignment_hooked_for() { - using mat = Eigen::Matrix; - using product = Eigen::Product; - namespace ei = Eigen::internal; - EXPECT_TRUE((std::is_base_of>, - ei::Assignment, - ei::Dense2Dense>>::value)); - EXPECT_TRUE( - (std::is_base_of>, - ei::Assignment, - ei::Dense2Dense>>::value)); - EXPECT_TRUE( - (std::is_base_of>, - ei::Assignment, - ei::Dense2Dense>>::value)); -} - -inline void expect_assignment_hooked() { - std::apply( - [](auto... t) { (expect_assignment_hooked_for(), ...); }, - ad_scalars{}); -} - inline Eigen::MatrixXd test_matrix(int rows, int cols, double offset) { Eigen::MatrixXd x(rows, cols); for (int j = 0; j < cols; ++j) { 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_product_hook_test.cpp b/test/unit/math/rev/core/eigen_product_hook_test.cpp index f4f0fa4176c..af9c1e90285 100644 --- a/test/unit/math/rev/core/eigen_product_hook_test.cpp +++ b/test/unit/math/rev/core/eigen_product_hook_test.cpp @@ -3,23 +3,21 @@ #include #include #include -#include -#include + +// Eigen's `operator*` on `var` matrices is evaluated by Stan's free functions +// (multiply, elt_multiply, scale_multiply, multiply_add). 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; -using stan::math::internal::is_stan_ad_product_hook_v; -using cvar = std::complex; template using mat = Eigen::Matrix; template using vec = Eigen::Matrix; template using row_vec = Eigen::Matrix; -template -using product_t = std::decay_t() - * std::declval())>; template mat make_mat(Eigen::Index rows, Eigen::Index cols, double offset) { @@ -32,61 +30,26 @@ mat make_mat(Eigen::Index rows, Eigen::Index cols, double offset) { return x; } -template -var sum_of(const Eigen::MatrixBase& x) { - var total = 0; - for (Eigen::Index j = 0; j < x.cols(); ++j) { - for (Eigen::Index i = 0; i < x.rows(); ++i) { - total += x.derived().coeff(i, j); - } - } - return total; +/** + * 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 stan::test::tape_delta direct_delta; + const auto expected = direct(); + const std::size_t direct_chain = direct_delta.chain(); + const std::size_t direct_nochain = direct_delta.nochain(); + const stan::test::tape_delta via_eigen_delta; + const auto actual = via_eigen(); + EXPECT_EQ(via_eigen_delta.chain(), direct_chain); + EXPECT_EQ(via_eigen_delta.nochain(), direct_nochain); + EXPECT_MATRIX_NEAR(stan::math::value_of(actual), + stan::math::value_of(expected), 1e-12); } } // namespace -TEST_F(AgradRev, eigen_product_hook_selects_var_products) { - EXPECT_TRUE((is_stan_ad_product_hook_v, mat>>)); - EXPECT_TRUE((is_stan_ad_product_hook_v, mat>>)); - EXPECT_TRUE((is_stan_ad_product_hook_v, mat>>)); - EXPECT_TRUE( - (is_stan_ad_product_hook_v, mat>>)); - EXPECT_TRUE(( - is_stan_ad_product_hook_v, mat>>)); - EXPECT_TRUE((is_stan_ad_product_hook_v, vec>>)); - EXPECT_TRUE((is_stan_ad_product_hook_v, vec>>)); - EXPECT_TRUE((is_stan_ad_product_hook_v, row_vec>>)); - EXPECT_TRUE((is_stan_ad_product_hook_v, mat>>)); - EXPECT_TRUE((is_stan_ad_product_hook_v< - product_t>, mat>>)); - EXPECT_TRUE( - (is_stan_ad_product_hook_v< - std::decay_t&>().lazyProduct( - std::declval&>()))>>)); - EXPECT_FALSE( - (is_stan_ad_product_hook_v, mat>>)); - EXPECT_FALSE((is_stan_ad_product_hook_v< - product_t, mat>>)); -} - -TEST_F(AgradRev, eigen_product_hook_selects_var_cwise_products) { - using arr = Eigen::Array; - using arr_d = Eigen::Array; - EXPECT_TRUE((is_stan_ad_product_hook_v>)); - EXPECT_TRUE((is_stan_ad_product_hook_v>)); - EXPECT_TRUE((is_stan_ad_product_hook_v>)); - EXPECT_TRUE((is_stan_ad_product_hook_v>>)); - EXPECT_TRUE((is_stan_ad_product_hook_v, var>>)); - EXPECT_TRUE((is_stan_ad_product_hook_v, int>>)); - EXPECT_TRUE((is_stan_ad_product_hook_v>>)); - EXPECT_TRUE( - (is_stan_ad_product_hook_v< - std::decay_t&>().cwiseProduct( - std::declval&>()))>>)); - EXPECT_FALSE((is_stan_ad_product_hook_v>)); - EXPECT_FALSE((is_stan_ad_product_hook_v>>)); -} - 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); @@ -105,236 +68,154 @@ TEST_F(AgradRev, eigen_product_hook_double_products_unchanged) { EXPECT_EQ(stan::test::tape_size(), 0); } -TEST_F(AgradRev, eigen_product_hook_matmul_one_callback) { - const Eigen::Index n = 5, k = 4, m = 3; - mat a = make_mat(n, k, 0.3); - const mat b = make_mat(k, m, -0.2); - stan::test::tape_delta delta; - mat c = a * b; - EXPECT_EQ(delta.chain(), 1); - EXPECT_EQ(delta.nochain(), n * m); - const mat expected = a.val() * b; - EXPECT_MATRIX_NEAR(c.val(), expected, 1e-12); - sum_of(c).grad(); - // d sum(A * B) / dA = ones(n, m) * B^T - const mat grad_a = mat::Ones(n, m) * b.transpose(); - EXPECT_MATRIX_NEAR(a.adj(), grad_a, 1e-12); +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([&] { return stan::math::multiply(a, b); }, + [&] { return mat(a.lazyProduct(b)); }); + 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_matmul_var_var) { - 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); - stan::test::tape_delta delta; - mat c = a * b; - EXPECT_EQ(delta.chain(), 1); - EXPECT_EQ(delta.nochain(), n * m); - sum_of(c).grad(); - EXPECT_MATRIX_NEAR(a.adj(), mat::Ones(n, m) * b.val().transpose(), - 1e-12); - EXPECT_MATRIX_NEAR(b.adj(), a.val().transpose() * mat::Ones(n, m), - 1e-12); +TEST_F(AgradRev, eigen_product_hook_fixed_size_products_use_multiply) { + mat f = make_mat(2, 2, 0.4); + const mat g = make_mat(2, 2, 0.1); + mat h = make_mat(4, 4, -0.3); + expect_same_as_direct([&] { return stan::math::multiply(f, g); }, + [&] { return mat(f * g); }); + expect_same_as_direct([&] { return stan::math::multiply(h, h); }, + [&] { return mat(h * h); }); } -TEST_F(AgradRev, eigen_product_hook_matmul_shapes) { - vec v = make_mat(4, 1, 0.1); +TEST_F(AgradRev, eigen_product_hook_inner_product_uses_multiply) { row_vec rv = make_mat(1, 4, 0.2); - const mat a = make_mat(3, 4, 0.3); - { - stan::test::tape_delta delta; - vec av = a * v; - EXPECT_EQ(delta.chain(), 1); - EXPECT_EQ(delta.nochain(), 3); - EXPECT_MATRIX_NEAR(av.val(), a * v.val(), 1e-12); - } - { - stan::test::tape_delta delta; - var dot = (rv * v)(0, 0); - EXPECT_EQ(delta.chain(), 1); - EXPECT_EQ(delta.nochain(), 1); - EXPECT_NEAR(dot.val(), rv.val().dot(v.val()), 1e-12); - } - { - stan::test::tape_delta delta; - mat outer = v * rv; - EXPECT_EQ(delta.chain(), 1); - EXPECT_EQ(delta.nochain(), 16); - EXPECT_MATRIX_NEAR(outer.val(), v.val() * rv.val(), 1e-12); - } - { - mat f = make_mat(2, 2, 0.4); - mat g = make_mat(2, 2, 0.1); - stan::test::tape_delta delta; - mat fg = f * g; - mat fg_lazy = f.lazyProduct(g); - EXPECT_EQ(delta.chain(), 2); - EXPECT_EQ(delta.nochain(), 8); - EXPECT_MATRIX_NEAR(fg.val(), f.val() * g, 1e-12); - EXPECT_MATRIX_NEAR(fg_lazy.val(), f.val() * g, 1e-12); - } -} - -TEST_F(AgradRev, eigen_product_hook_assignment_forms) { - mat a = make_mat(3, 3, 0.3); - mat b = make_mat(3, 3, -0.1); - mat c = make_mat(3, 3, 0.05); - const mat c0 = c.val(); - { - stan::test::tape_delta delta; - mat d(3, 3); - d.noalias() = a * b; - EXPECT_EQ(delta.chain(), 1); - EXPECT_EQ(delta.nochain(), 9); - } - { - stan::test::tape_delta delta; - c += a * b; - // one callback for the product plus one vari per coefficient for += - EXPECT_EQ(delta.chain(), 1 + 9); - EXPECT_EQ(delta.nochain(), 9); - EXPECT_MATRIX_NEAR(c.val(), c0 + a.val() * b.val(), 1e-12); - } - { - // aliasing: the product is evaluated before the destination is written - mat e = a; - e = e * b; - EXPECT_MATRIX_NEAR(e.val(), a.val() * b.val(), 1e-12); - } + vec v = make_mat(4, 1, 0.1); + const stan::test::tape_delta direct_delta; + const var expected = stan::math::multiply(rv, v); + const std::size_t direct_total = direct_delta.total(); + const stan::test::tape_delta via_eigen_delta; + const var actual = (rv * v)(0, 0); + EXPECT_EQ(via_eigen_delta.total(), direct_total); + EXPECT_NEAR(actual.val(), expected.val(), 1e-12); } -TEST_F(AgradRev, eigen_product_hook_cwise_one_callback) { - const Eigen::Index n = 4; - vec x = make_mat(n, 1, 0.3); - vec y = make_mat(n, 1, -0.4); - stan::test::tape_delta delta; - vec z = (x.array() * y.array()).matrix(); - EXPECT_EQ(delta.chain(), 1); - EXPECT_EQ(delta.nochain(), n); - sum_of(z).grad(); - EXPECT_MATRIX_NEAR(x.adj(), y.val(), 1e-12); - EXPECT_MATRIX_NEAR(y.adj(), x.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_times_matrix) { +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; - { - stan::test::tape_delta delta; - mat c = s * a; - EXPECT_EQ(delta.chain(), 1); - EXPECT_EQ(delta.nochain(), 6); - EXPECT_MATRIX_NEAR(c.val(), 1.5 * a, 1e-12); - sum_of(c).grad(); - EXPECT_NEAR(s.adj(), a.sum(), 1e-12); - } - stan::math::set_zero_all_adjoints(); - mat av = a; - { - stan::test::tape_delta delta; - mat c = av * 2; - EXPECT_EQ(delta.chain(), 1); - EXPECT_EQ(delta.nochain(), 6); - sum_of(c).grad(); - EXPECT_MATRIX_NEAR(av.adj(), mat::Constant(3, 2, 2.0), 1e-12); - } + 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_fuse) { - const Eigen::Index n = 3, k = 4, m = 2; - const mat a = make_mat(n, k, 0.1); - mat b = make_mat(k, m, -0.3); +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; - const mat expected = 0.7 * (a * b.val()); - { - stan::test::tape_delta delta; - mat c = s * (a * b); - EXPECT_EQ(delta.chain(), 1); - EXPECT_EQ(delta.nochain(), n * m); - EXPECT_MATRIX_NEAR(c.val(), expected, 1e-12); - sum_of(c).grad(); - EXPECT_NEAR(s.adj(), (a * b.val()).sum(), 1e-12); - EXPECT_MATRIX_NEAR(b.adj(), 0.7 * a.transpose() * mat::Ones(n, m), - 1e-12); - } - { - stan::test::tape_delta delta; - mat c = (a * b) * s; - EXPECT_EQ(delta.chain(), 1); - EXPECT_EQ(delta.nochain(), n * m); - EXPECT_MATRIX_NEAR(c.val(), expected, 1e-12); - } - { - stan::test::tape_delta delta; - mat c = (s * a) * b; - EXPECT_EQ(delta.chain(), 1); - EXPECT_EQ(delta.nochain(), n * m); - EXPECT_MATRIX_NEAR(c.val(), expected, 1e-12); - } - { - stan::test::tape_delta delta; - mat c = a * (b * s); - EXPECT_EQ(delta.chain(), 1); - EXPECT_EQ(delta.nochain(), n * m); - EXPECT_MATRIX_NEAR(c.val(), expected, 1e-12); - } - { - mat d = make_mat(n, m, 0.2); - stan::test::tape_delta delta; - mat c = d + s * (a * b); - // one product callback plus one vari per coefficient for the sum - EXPECT_EQ(delta.chain(), 1 + n * m); - EXPECT_EQ(delta.nochain(), n * m); - EXPECT_MATRIX_NEAR(c.val(), d.val() + expected, 1e-12); - } + 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_scales_are_small) { - // every product callback captures its scalar factors by value, so they - // must not embed storage for the factors themselves - EXPECT_LE(sizeof(stan::math::internal::rev_product_scales), - 4 * sizeof(void*)); - EXPECT_LE( - sizeof(stan::math::internal::rev_product_scales>), - 5 * sizeof(void*)); +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::subtract(c, stan::math::multiply(a, b)); }, + [&] { + 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_many_autodiff_scales) { - const mat a = make_mat(2, 3, 0.1); - const mat b = make_mat(3, 2, -0.2); - var s1 = 0.5, s2 = 2.0, s3 = -1.5; - mat c = s1 * (s2 * a) * (b * s3); - const double scale = 0.5 * 2.0 * -1.5; - EXPECT_MATRIX_NEAR(c.val(), scale * (a * b), 1e-12); - sum_of(c).grad(); - const double total = (a * b).sum(); - EXPECT_NEAR(s1.adj(), 2.0 * -1.5 * total, 1e-12); - EXPECT_NEAR(s2.adj(), 0.5 * -1.5 * total, 1e-12); - EXPECT_NEAR(s3.adj(), 0.5 * 2.0 * total, 1e-12); +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_complex_matmul) { - mat a(2, 2); - a << cvar(1.0, 0.5), cvar(-0.3, 0.2), cvar(0.4, -1.0), cvar(2.0, 0.1); - mat> b(2, 1); - b << std::complex(0.5, -0.25), std::complex(1.5, 0.75); - stan::test::tape_delta delta; - mat c = a * b; - EXPECT_EQ(delta.chain(), 1); - EXPECT_EQ(delta.nochain(), 4); - mat> a_val(2, 2); - for (Eigen::Index j = 0; j < 2; ++j) { - for (Eigen::Index i = 0; i < 2; ++i) { - a_val(i, j) = {a(i, j).real().val(), a(i, j).imag().val()}; - } - } - const mat> expected = a_val * b; - for (Eigen::Index i = 0; i < 2; ++i) { - EXPECT_NEAR(c(i, 0).real().val(), expected(i, 0).real(), 1e-12); - EXPECT_NEAR(c(i, 0).imag().val(), expected(i, 0).imag(), 1e-12); - } - // d re(c_0) / d re(a_00) = re(b_0), d re(c_0) / d im(a_00) = -im(b_0) - c(0, 0).real().grad(); - EXPECT_NEAR(a(0, 0).real().adj(), b(0).real(), 1e-12); - EXPECT_NEAR(a(0, 0).imag().adj(), -b(0).imag(), 1e-12); - EXPECT_NEAR(a(1, 0).real().adj(), 0.0, 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); } diff --git a/test/unit/math/rev/fun/multiply_add_test.cpp b/test/unit/math/rev/fun/multiply_add_test.cpp new file mode 100644 index 00000000000..051863ecc82 --- /dev/null +++ b/test/unit/math/rev/fun/multiply_add_test.cpp @@ -0,0 +1,48 @@ +#include +#include +#include +#include +#include + +namespace { +Eigen::MatrixXd multiply_add_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 + +TEST_F(AgradRev, multiply_add_values_and_gradients) { + using stan::math::var; + const int n = 3, k = 4, m = 2; + Eigen::Matrix a = multiply_add_test_matrix(n, k, 0.1); + Eigen::Matrix b = multiply_add_test_matrix(k, m, -0.2); + Eigen::Matrix c = multiply_add_test_matrix(n, m, 0.4); + const stan::test::tape_delta delta; + Eigen::Matrix d = stan::math::multiply_add(a, b, c); + // one callback plus one vari per result coefficient + EXPECT_EQ(delta.total(), 1 + n * m); + EXPECT_MATRIX_NEAR(d.val(), a.val() * b.val() + c.val(), 1e-12); + stan::math::sum(d).grad(); + const Eigen::MatrixXd ones = Eigen::MatrixXd::Ones(n, m); + EXPECT_MATRIX_NEAR(a.adj(), ones * b.val().transpose(), 1e-12); + EXPECT_MATRIX_NEAR(b.adj(), a.val().transpose() * ones, 1e-12); + EXPECT_MATRIX_NEAR(c.adj(), ones, 1e-12); +} + +TEST_F(AgradRev, multiply_add_double_addend) { + using stan::math::var; + const int n = 2, k = 3, m = 2; + Eigen::Matrix a = multiply_add_test_matrix(n, k, 0.1); + const Eigen::MatrixXd b = multiply_add_test_matrix(k, m, -0.2); + const Eigen::MatrixXd c = multiply_add_test_matrix(n, m, 0.4); + Eigen::Matrix d = stan::math::multiply_add(a, b, c); + EXPECT_MATRIX_NEAR(d.val(), a.val() * b + c, 1e-12); + stan::math::sum(d).grad(); + EXPECT_MATRIX_NEAR(a.adj(), Eigen::MatrixXd::Ones(n, m) * b.transpose(), + 1e-12); +} diff --git a/test/unit/math/rev/fun/scale_multiply_test.cpp b/test/unit/math/rev/fun/scale_multiply_test.cpp new file mode 100644 index 00000000000..34f6e6403f6 --- /dev/null +++ b/test/unit/math/rev/fun/scale_multiply_test.cpp @@ -0,0 +1,47 @@ +#include +#include +#include +#include +#include + +namespace { +Eigen::MatrixXd scale_multiply_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 + +TEST_F(AgradRev, scale_multiply_values_and_gradients) { + using stan::math::var; + const int n = 3, k = 4, m = 2; + var s = 0.7; + Eigen::Matrix a = scale_multiply_test_matrix(n, k, 0.1); + Eigen::Matrix b = scale_multiply_test_matrix(k, m, -0.2); + const stan::test::tape_delta delta; + Eigen::Matrix c = stan::math::scale_multiply(s, a, b); + // one callback plus one vari per result coefficient + EXPECT_EQ(delta.total(), 1 + n * m); + EXPECT_MATRIX_NEAR(c.val(), 0.7 * a.val() * b.val(), 1e-12); + stan::math::sum(c).grad(); + const Eigen::MatrixXd ones = Eigen::MatrixXd::Ones(n, m); + EXPECT_NEAR(s.adj(), (a.val() * b.val()).sum(), 1e-12); + EXPECT_MATRIX_NEAR(a.adj(), 0.7 * ones * b.val().transpose(), 1e-12); + EXPECT_MATRIX_NEAR(b.adj(), 0.7 * a.val().transpose() * ones, 1e-12); +} + +TEST_F(AgradRev, scale_multiply_double_scale) { + using stan::math::var; + const int n = 2, k = 3, m = 2; + const Eigen::MatrixXd a = scale_multiply_test_matrix(n, k, 0.1); + Eigen::Matrix b = scale_multiply_test_matrix(k, m, -0.2); + Eigen::Matrix c = stan::math::scale_multiply(-1.5, a, b); + EXPECT_MATRIX_NEAR(c.val(), -1.5 * a * b.val(), 1e-12); + stan::math::sum(c).grad(); + EXPECT_MATRIX_NEAR(b.adj(), + -1.5 * a.transpose() * Eigen::MatrixXd::Ones(n, m), 1e-12); +} From 9a6a82bad8a9c21fa988f8a0fd4d76e64770946f Mon Sep 17 00:00:00 2001 From: Steve Bronder Date: Tue, 13 Oct 2026 16:16:04 -0400 Subject: [PATCH 4/7] Simplify the Eigen specializations that route var products to Stan Move the hooks header to stan/math/rev/meta and replace the macros with plain specializations of Eigen's structs: - binary_evaluator: one specialization. Its var condition sits in the CwiseBinaryOp functor argument, which Eigen's specialization leaves free, so it covers var with var, double and int in one place. - Assignment: one specialization per functor (=, +=, -=). The condition sits in both scalars of the functor, so products of non-dense operands run Eigen's own specialization and the copied Eigen bodies are gone. Inner products are now computed by multiply as well. - product_evaluator and binary_evaluator share one evaluator body, ad_product_evaluator, which no longer allocates a result it discards. - Lazy products (lazyProduct()) are left to Eigen. The traits strip references and const themselves, with require_eigen_product_t and require_eigen_cwise_product_t aliases. Every Stan specialization derives from eigen_ad_hook so tests can check which struct Eigen picks. scale_multiply and multiply_add return their arena result instead of copying it to the heap. Tests: eigen_hook_dispatch_test checks which specialization Eigen picks for each scalar pairing, operator and destination; MRE tests instantiate each struct on its own; eigen_product_hook_test gains products nested in larger expressions, Stan functions and block/transpose destinations. --- stan/math/prim/fun/Eigen.hpp | 2 - .../meta/eigen_product_evaluator_hooks.hpp | 378 ---------------- stan/math/rev/core/Eigen_NumTraits.hpp | 24 +- stan/math/rev/fun/multiply_add.hpp | 2 +- stan/math/rev/fun/scale_multiply.hpp | 2 +- stan/math/rev/meta.hpp | 1 + .../meta/eigen_product_evaluator_hooks.hpp | 415 ++++++++++++++++++ .../mix/fun/multiply_operator_matmul_test.cpp | 18 +- ...eigen_cwise_product_evaluator_mre_test.cpp | 15 + .../eigen_lazy_product_evaluator_mre_test.cpp | 15 + .../eigen_product_assignment_mre_test.cpp | 16 + .../core/eigen_product_evaluator_mre_test.cpp | 14 + .../math/rev/core/eigen_product_hook_test.cpp | 167 ++++++- .../math/rev/fun/eigen_hook_dispatch_test.cpp | 159 +++++++ 14 files changed, 820 insertions(+), 408 deletions(-) delete mode 100644 stan/math/prim/meta/eigen_product_evaluator_hooks.hpp create mode 100644 stan/math/rev/meta/eigen_product_evaluator_hooks.hpp create mode 100644 test/unit/math/rev/core/eigen_cwise_product_evaluator_mre_test.cpp create mode 100644 test/unit/math/rev/core/eigen_lazy_product_evaluator_mre_test.cpp create mode 100644 test/unit/math/rev/core/eigen_product_assignment_mre_test.cpp create mode 100644 test/unit/math/rev/core/eigen_product_evaluator_mre_test.cpp create mode 100644 test/unit/math/rev/fun/eigen_hook_dispatch_test.cpp diff --git a/stan/math/prim/fun/Eigen.hpp b/stan/math/prim/fun/Eigen.hpp index 47f8d9867d0..9678be671ad 100644 --- a/stan/math/prim/fun/Eigen.hpp +++ b/stan/math/prim/fun/Eigen.hpp @@ -87,6 +87,4 @@ } // namespace Eigen -#include - #endif diff --git a/stan/math/prim/meta/eigen_product_evaluator_hooks.hpp b/stan/math/prim/meta/eigen_product_evaluator_hooks.hpp deleted file mode 100644 index 117bc61d802..00000000000 --- a/stan/math/prim/meta/eigen_product_evaluator_hooks.hpp +++ /dev/null @@ -1,378 +0,0 @@ -#ifndef STAN_MATH_PRIM_META_EIGEN_PRODUCT_EVALUATOR_HOOKS_HPP -#define STAN_MATH_PRIM_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` and `multiply_add`, so that `A * B` records one node on - * the reverse-mode stack instead of one per scalar operation. - * - * Eigen's `product_evaluator` takes the operand scalar types as defaulted - * template parameters, so a specialization that fixes them is ambiguous with - * Eigen's generic one. The dense `product_evaluator` below therefore matches - * every dense-by-dense product and handles non-autodiff scalars exactly as - * Eigen's generic `product_evaluator` in Eigen/src/Core/ProductEvaluators.h - * does. When upgrading Eigen, compare that evaluator with this one. - * - * Because the specializations match every dense product, this file is - * included from stan/math/prim/fun/Eigen.hpp so that every translation unit - * sees them. The functions they 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 - -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::Scalar>::value - || is_var::Scalar>::value) - && !is_complex::Scalar>::value - && !is_complex::Scalar>::value; - -/** - * 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 - */ -template -struct is_scalar_times_matrix : std::false_type {}; -template -struct is_scalar_times_matrix, - const Eigen::CwiseNullaryOp, Plain>, - const Mat>> : 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 - */ -template -struct is_matrix_times_scalar : std::false_type {}; -template -struct is_matrix_times_scalar< - Eigen::CwiseBinaryOp, const Mat, - const Eigen::CwiseNullaryOp< - Eigen::internal::scalar_constant_op, Plain>>> - : std::true_type {}; - -/** - * True for an Eigen matrix product expression. - * - * @tparam T an Eigen expression type - */ -template -struct is_eigen_product : std::false_type {}; -template -struct is_eigen_product> : std::true_type {}; - -} // namespace internal -} // namespace math -} // namespace stan - -namespace Eigen { -namespace internal { - -/** - * Dense-by-dense products. For operands with `var` scalars the product is - * `scale_multiply(s, A, B)` when either operand is Eigen's `s * A` (Eigen - * rewrites `s * (A * B)` as `(s * A) * B`) and `multiply(lhs, rhs)` - * otherwise. For all other scalars this is Eigen's generic - * `product_evaluator`. More specialized than Eigen's generic - * `product_evaluator` (fixed shapes) and less specialized than its lazy and - * `LazyCoeffBasedProductMode` evaluators, which are handled below. - */ -template -struct product_evaluator, ProductTag, DenseShape, - DenseShape> - : public evaluator::PlainObject> { - using XprType = Product; - using PlainObject = typename XprType::PlainObject; - using Base = evaluator; - enum { Flags = Base::Flags | EvalBeforeNestingBit }; - - explicit product_evaluator(const XprType& xpr) - : m_result(xpr.rows(), xpr.cols()) { - using stan::math::internal::is_matrix_times_scalar; - using stan::math::internal::is_scalar_times_matrix; - if constexpr (!stan::math::internal::is_stan_ad_product_v) { - generic_product_impl::evalTo(m_result, xpr.lhs(), xpr.rhs()); - } else if constexpr (is_scalar_times_matrix::value) { - // (s * A) * B - m_result = scale_multiply(xpr.lhs().lhs().functor().m_other, - xpr.lhs().rhs(), xpr.rhs()); - } else if constexpr (is_matrix_times_scalar::value) { - // (A * s) * B - m_result = scale_multiply(xpr.lhs().rhs().functor().m_other, - xpr.lhs().lhs(), xpr.rhs()); - } else if constexpr (is_scalar_times_matrix::value) { - // A * (s * B) - m_result = scale_multiply(xpr.rhs().lhs().functor().m_other, xpr.lhs(), - xpr.rhs().rhs()); - } else if constexpr (is_matrix_times_scalar::value) { - // A * (B * s) - m_result = scale_multiply(xpr.rhs().rhs().functor().m_other, xpr.lhs(), - xpr.rhs().lhs()); - } else if constexpr (stan::is_stan_scalar::value) { - // a row vector times a column vector - m_result.coeffRef(0, 0) = multiply(xpr.lhs(), xpr.rhs()); - } else { - m_result = multiply(xpr.lhs(), xpr.rhs()); - } - // 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. - construct_at(this, m_result); - } - - protected: - PlainObject m_result; -}; - -/** - * Lazy (coefficient-based) products, one specialization per product tag so - * each is more specialized than Eigen's `LazyProduct` evaluator. For `var` - * operands they use the dense evaluator above (the product tag is unused - * there); for other scalars Eigen's lazy evaluator, reached with tag -1, - * which no product has and which Eigen's lazy evaluator ignores. Eigen's - * `LazyCoeffBasedProductMode` evaluator for default products derives from - * the `CoeffBasedProductMode` specialization, so it is covered as well. - */ -#define STAN_EIGEN_LAZY_PRODUCT_EVALUATOR(TAG) \ - template \ - struct product_evaluator, TAG, DenseShape, \ - DenseShape> \ - : std::conditional_t< \ - stan::math::internal::is_stan_ad_product_v, \ - product_evaluator, GemmProduct, \ - DenseShape, DenseShape>, \ - product_evaluator, -1, DenseShape, \ - DenseShape>> { \ - using XprType = Product; \ - using Base = std::conditional_t< \ - stan::math::internal::is_stan_ad_product_v, \ - product_evaluator, GemmProduct, \ - DenseShape, DenseShape>, \ - product_evaluator, -1, DenseShape, \ - DenseShape>>; \ - explicit product_evaluator(const XprType& xpr) : Base(base_xpr(xpr)) {} \ - \ - private: \ - static auto base_xpr(const XprType& xpr) { \ - if constexpr (stan::math::internal::is_stan_ad_product_v) { \ - return Product(xpr.lhs(), xpr.rhs()); \ - } else { \ - return xpr; \ - } \ - } \ - }; - -STAN_EIGEN_LAZY_PRODUCT_EVALUATOR(InnerProduct) -STAN_EIGEN_LAZY_PRODUCT_EVALUATOR(OuterProduct) -STAN_EIGEN_LAZY_PRODUCT_EVALUATOR(CoeffBasedProductMode) -STAN_EIGEN_LAZY_PRODUCT_EVALUATOR(LazyCoeffBasedProductMode) -STAN_EIGEN_LAZY_PRODUCT_EVALUATOR(GemvProduct) -STAN_EIGEN_LAZY_PRODUCT_EVALUATOR(GemmProduct) -#undef STAN_EIGEN_LAZY_PRODUCT_EVALUATOR - -} // namespace internal -} // namespace Eigen - -/** - * Specialize Eigen's `binary_evaluator` for the coefficient-wise product - * `lhs .* rhs` with scalar pattern `(LS, RS)`. The product is - * `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. `TPL` is a - * function-like macro expanding to extra template parameters with a - * trailing comma. The scalar template parameters of `binary_evaluator` stay - * defaulted, as in Eigen's generic specialization, so this one is more - * specialized only through the functor. Use inside `namespace - * Eigen::internal`. - */ -#define STAN_EIGEN_AD_CWISE_PRODUCT_EVALUATOR(TPL, LS, RS) \ - template \ - struct binary_evaluator, Lhs, Rhs>, \ - IndexBased, IndexBased> \ - : public evaluator, \ - Lhs, Rhs>::PlainObject> { \ - using XprType = CwiseBinaryOp, Lhs, Rhs>; \ - using PlainObject = typename XprType::PlainObject; \ - using Base = evaluator; \ - enum { Flags = Base::Flags | EvalBeforeNestingBit }; \ - \ - explicit binary_evaluator(const XprType& xpr) \ - : m_result(xpr.rows(), xpr.cols()) { \ - using stan::math::internal::is_eigen_product; \ - using stan::math::internal::is_matrix_times_scalar; \ - using stan::math::internal::is_scalar_times_matrix; \ - /* Stan's functions take matrices, so array operands are viewed as */ \ - /* matrices (`.matrix().derived()` is the operand itself for a */ \ - /* matrix) and the result is viewed back as an array */ \ - const auto product = [&xpr]() { \ - if constexpr (is_scalar_times_matrix::value) { \ - if constexpr (is_eigen_product>::value) { \ - /* s * (A * B) */ \ - return scale_multiply(xpr.lhs().functor().m_other, \ - xpr.rhs().lhs(), xpr.rhs().rhs()); \ - } else { \ - /* s * X */ \ - return multiply(xpr.lhs().functor().m_other, \ - xpr.rhs().matrix().derived()); \ - } \ - } else if constexpr (is_matrix_times_scalar::value) { \ - if constexpr (is_eigen_product>::value) { \ - /* (A * B) * s */ \ - return scale_multiply(xpr.rhs().functor().m_other, \ - xpr.lhs().lhs(), xpr.lhs().rhs()); \ - } else { \ - /* X * s */ \ - return multiply(xpr.rhs().functor().m_other, \ - xpr.lhs().matrix().derived()); \ - } \ - } else { \ - return elt_multiply(xpr.lhs().matrix().derived(), \ - xpr.rhs().matrix().derived()); \ - } \ - }(); \ - if constexpr (std::is_base_of, \ - PlainObject>::value) { \ - m_result = product.array(); \ - } else { \ - m_result = product; \ - } \ - /* Base points into m_result's storage, so it is constructed only */ \ - /* once m_result is filled */ \ - construct_at(this, m_result); \ - } \ - \ - protected: \ - PlainObject m_result; \ - }; - -/** - * Specialize Eigen's `Assignment` of a product to a destination with scalar - * `SCALAR` for `OP` (`assign_op`, `add_assign_op` or `sub_assign_op`). - * Eigen's own specializations for this case compute the product directly - * instead of through `product_evaluator`. For dense operands, `dst = A * B` - * uses `multiply` (or `scale_multiply` for a scaled operand), `dst += A * B` - * uses `multiply_add(A, B, dst)` and `dst -= A * B` uses `subtract(dst, - * multiply(A, B))`; other shapes (e.g. triangular views) run Eigen's body. - * The last template argument repeats Eigen's own condition token for token, - * so this is more specialized only through the functor's scalar. Use inside - * `namespace Eigen::internal`. - */ -#define STAN_EIGEN_AD_PRODUCT_ASSIGNMENT_OP(TPL, SCALAR, OP) \ - template \ - struct Assignment, \ - OP, Dense2Dense, \ - std::enable_if_t<(Options == DefaultProduct \ - || Options == AliasFreeProduct)>> { \ - using SrcXprType = Product; \ - static void run(DstXprType& dst, const SrcXprType& src, \ - const OP& func) { \ - using stan::math::internal::is_matrix_times_scalar; \ - using stan::math::internal::is_scalar_times_matrix; \ - constexpr bool is_add \ - = std::is_same, \ - add_assign_op>::value; \ - constexpr bool is_sub \ - = std::is_same, \ - sub_assign_op>::value; \ - constexpr bool is_scaled = is_scalar_times_matrix::value \ - || is_matrix_times_scalar::value \ - || is_scalar_times_matrix::value \ - || is_matrix_times_scalar::value; \ - if constexpr (!std::is_same::Shape, \ - DenseShape>::value \ - || !std::is_same::Shape, \ - DenseShape>::value) { \ - if constexpr (is_add) { \ - generic_product_impl::addTo(dst, src.lhs(), src.rhs()); \ - } else if constexpr (is_sub) { \ - generic_product_impl::subTo(dst, src.lhs(), src.rhs()); \ - } else { \ - if (dst.rows() != src.rows() || dst.cols() != src.cols()) { \ - dst.resize(src.rows(), src.cols()); \ - } \ - generic_product_impl::evalTo(dst, src.lhs(), src.rhs()); \ - } \ - } else 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); \ - } \ - func.assignCoeff(dst.coeffRef(0, 0), multiply(src.lhs(), src.rhs())); \ - } else if constexpr (is_add && !is_scaled) { \ - dst = multiply_add(src.lhs(), src.rhs(), dst); \ - } else if constexpr (is_sub && !is_scaled) { \ - dst = subtract(dst, multiply(src.lhs(), src.rhs())); \ - } else { \ - const auto product = [&src]() { \ - if constexpr (is_scalar_times_matrix::value) { \ - /* (s * A) * B */ \ - return scale_multiply(src.lhs().lhs().functor().m_other, \ - src.lhs().rhs(), src.rhs()); \ - } else if constexpr (is_matrix_times_scalar::value) { \ - /* (A * s) * B */ \ - return scale_multiply(src.lhs().rhs().functor().m_other, \ - src.lhs().lhs(), src.rhs()); \ - } else if constexpr (is_scalar_times_matrix::value) { \ - /* A * (s * B) */ \ - return scale_multiply(src.rhs().lhs().functor().m_other, \ - src.lhs(), src.rhs().rhs()); \ - } else if constexpr (is_matrix_times_scalar::value) { \ - /* A * (B * s) */ \ - return scale_multiply(src.rhs().rhs().functor().m_other, \ - src.lhs(), src.rhs().lhs()); \ - } else { \ - return multiply(src.lhs(), src.rhs()); \ - } \ - }(); \ - if constexpr (is_add) { \ - dst += product; \ - } else if constexpr (is_sub) { \ - dst -= product; \ - } else { \ - dst = product; \ - } \ - } \ - } \ - }; - -#define STAN_EIGEN_AD_PRODUCT_ASSIGNMENT(TPL, SCALAR) \ - STAN_EIGEN_AD_PRODUCT_ASSIGNMENT_OP(TPL, SCALAR, assign_op) \ - STAN_EIGEN_AD_PRODUCT_ASSIGNMENT_OP(TPL, SCALAR, add_assign_op) \ - STAN_EIGEN_AD_PRODUCT_ASSIGNMENT_OP(TPL, SCALAR, sub_assign_op) - -#endif diff --git a/stan/math/rev/core/Eigen_NumTraits.hpp b/stan/math/rev/core/Eigen_NumTraits.hpp index 6ca7a84c293..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 @@ -302,29 +303,6 @@ struct ternary_evaluator< }; #endif -/** - * Coefficient-wise products and products with a scalar of `var` with `var`, - * `double` or `int` are computed by `elt_multiply`, `multiply` or - * `scale_multiply` (see stan/math/prim/meta/eigen_product_evaluator_hooks.hpp). - */ -#define STAN_REV_NO_TPARAMS() -STAN_EIGEN_AD_CWISE_PRODUCT_EVALUATOR(STAN_REV_NO_TPARAMS, stan::math::var, - stan::math::var) -STAN_EIGEN_AD_CWISE_PRODUCT_EVALUATOR(STAN_REV_NO_TPARAMS, stan::math::var, - double) -STAN_EIGEN_AD_CWISE_PRODUCT_EVALUATOR(STAN_REV_NO_TPARAMS, double, - stan::math::var) -STAN_EIGEN_AD_CWISE_PRODUCT_EVALUATOR(STAN_REV_NO_TPARAMS, stan::math::var, int) -STAN_EIGEN_AD_CWISE_PRODUCT_EVALUATOR(STAN_REV_NO_TPARAMS, int, stan::math::var) - -/** - * Products assigned to (`=`, `+=`, `-=`) a destination with `var` scalars are - * computed by `multiply`, `scale_multiply`, `multiply_add` or `subtract` (see - * stan/math/prim/meta/eigen_product_evaluator_hooks.hpp). - */ -STAN_EIGEN_AD_PRODUCT_ASSIGNMENT(STAN_REV_NO_TPARAMS, stan::math::var) -#undef STAN_REV_NO_TPARAMS - /** * Enable linear access of inputs when using read_vi_val_adj. */ diff --git a/stan/math/rev/fun/multiply_add.hpp b/stan/math/rev/fun/multiply_add.hpp index 8a4d310d21e..3b6f6b08ef8 100644 --- a/stan/math/rev/fun/multiply_add.hpp +++ b/stan/math/rev/fun/multiply_add.hpp @@ -60,7 +60,7 @@ inline auto multiply_add(const Mat1& A, const Mat2& B, const Mat3& C) { arena_C.adj() += res_adj; } }); - return return_t(res); + return res; } } // namespace math diff --git a/stan/math/rev/fun/scale_multiply.hpp b/stan/math/rev/fun/scale_multiply.hpp index a7f45add5fa..9a5e101b9f5 100644 --- a/stan/math/rev/fun/scale_multiply.hpp +++ b/stan/math/rev/fun/scale_multiply.hpp @@ -57,7 +57,7 @@ inline auto scale_multiply(const Scal& s, const Mat1& A, const Mat2& B) { arena_B.adj() += s_val * arena_A_val.transpose() * res_adj; } }); - return return_t(res); + return res; } } // namespace math 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..33c5438bd3f --- /dev/null +++ b/stan/math/rev/meta/eigen_product_evaluator_hooks.hpp @@ -0,0 +1,415 @@ +#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` and `multiply_add`, 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 + +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 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, std::forward(lhs).rhs(), + std::forward(rhs)); + } else if constexpr (is_matrix_times_scalar::value) { + // (A * s) * B + return scale_multiply(std::forward(lhs).rhs().functor().m_other, std::forward(lhs).lhs(), + std::forward(rhs)); + } else if constexpr (is_scalar_times_matrix::value) { + // A * (s * B) + return scale_multiply(std::forward(rhs).lhs().functor().m_other, std::forward(lhs), + 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, std::forward(lhs), + std::forward(rhs).lhs()); + } else { + return multiply(std::forward(lhs), 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, std::forward(xpr).rhs().lhs(), + std::forward(xpr).rhs().rhs()); + } else { + // s * X + return multiply(std::forward(xpr).lhs().functor().m_other, + 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, + std::forward(xpr).lhs().lhs(), + std::forward(xpr).lhs().rhs()); + } else { + // X * s + return multiply(std::forward(xpr).rhs().functor().m_other, + std::forward(xpr).lhs().matrix().derived()); + } + } else { + return elt_multiply(std::forward(xpr).lhs().matrix().derived(), + 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::value) { + // 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 (std::is_base_of, + PlainObject>::value) { + 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 */) { + dst = multiply_add(std::forward(src).lhs(), std::forward(src).rhs(), 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 */) { + dst = subtract( + dst, stan::math::internal::multiply_operands(std::forward(src).lhs(), std::forward(src).rhs())); + } +}; + +} // namespace internal +} // namespace Eigen + +#endif diff --git a/test/unit/math/mix/fun/multiply_operator_matmul_test.cpp b/test/unit/math/mix/fun/multiply_operator_matmul_test.cpp index 2314f91200d..db7562fc036 100644 --- a/test/unit/math/mix/fun/multiply_operator_matmul_test.cpp +++ b/test/unit/math/mix/fun/multiply_operator_matmul_test.cpp @@ -42,7 +42,7 @@ TEST(mathMixMultiplyOperator, innerAndOuterProduct) { } TEST(mathMixMultiplyOperator, fixedSize) { - // 2x2 uses Eigen's lazy coefficient-based product + // 2x2 uses Eigen's coefficient-based product auto f_fixed = [](const auto& a, const auto& b) { using a_scalar = typename std::decay_t::Scalar; using b_scalar = typename std::decay_t::Scalar; @@ -56,7 +56,23 @@ TEST(mathMixMultiplyOperator, fixedSize) { test_matrix(2, 2, -0.3)); } +TEST(mathMixMultiplyOperator, fixedSizeOuterProduct) { + // a fixed-size outer product uses Eigen's lazy coefficient-based product + auto f_outer = [](const auto& a, const auto& b) { + using a_scalar = typename std::decay_t::Scalar; + using b_scalar = typename std::decay_t::Scalar; + using ret_scalar = stan::return_type_t; + const Eigen::Matrix a3 = a; + const Eigen::Matrix b3 = b; + const Eigen::Matrix c = a3 * b3; + return Eigen::Matrix(c); + }; + stan::test::expect_ad(f_outer, test_matrix(3, 1, 0.1), + test_matrix(1, 3, -0.3)); +} + 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); 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 index af9c1e90285..400451464aa 100644 --- a/test/unit/math/rev/core/eigen_product_hook_test.cpp +++ b/test/unit/math/rev/core/eigen_product_hook_test.cpp @@ -85,8 +85,6 @@ TEST_F(AgradRev, eigen_product_hook_matrix_products_use_multiply) { [&] { return vec(a * v); }); expect_same_as_direct([&] { return stan::math::multiply(v, rv); }, [&] { return mat(v * rv); }); - expect_same_as_direct([&] { return stan::math::multiply(a, b); }, - [&] { return mat(a.lazyProduct(b)); }); expect_same_as_direct( [&] { mat c(5, 3); @@ -219,3 +217,168 @@ TEST_F(AgradRev, eigen_product_hook_gradients) { 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); +} 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..6a8fa3a5f95 --- /dev/null +++ b/test/unit/math/rev/fun/eigen_hook_dispatch_test.cpp @@ -0,0 +1,159 @@ +#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