Override Eigen's product evaluators

stan/math/fwd/fun/Eigen_NumTraits.hpp

@@ -253,4 +253,5 @@ struct ternary_evaluator<
 }  // namespace internal
 
 }  // namespace Eigen
+
 #endif

stan/math/prim/fun.hpp

@@ -223,8 +223,10 @@
 #include <stan/math/prim/fun/modified_bessel_second_kind.hpp>
 #include <stan/math/prim/fun/modulus.hpp>
 #include <stan/math/prim/fun/multiply.hpp>
+#include <stan/math/prim/fun/multiply_add.hpp>
 #include <stan/math/prim/fun/multiply_log.hpp>
 #include <stan/math/prim/fun/multiply_lower_tri_self_transpose.hpp>
+#include <stan/math/prim/fun/multiply_sub.hpp>
 #include <stan/math/prim/fun/norm.hpp>
 #include <stan/math/prim/fun/norm1.hpp>
 #include <stan/math/prim/fun/norm2.hpp>
@@ -281,6 +283,7 @@
 #include <stan/math/prim/fun/rows_dot_self.hpp>
 #include <stan/math/prim/fun/scalbn.hpp>
 #include <stan/math/prim/fun/scale_matrix_exp_multiply.hpp>
+#include <stan/math/prim/fun/scale_multiply.hpp>
 #include <stan/math/prim/fun/scalar_seq_view.hpp>
 #include <stan/math/prim/fun/scaled_add.hpp>
 #include <stan/math/prim/fun/sd.hpp>

stan/math/prim/fun/multiply_add.hpp

@@ -0,0 +1,46 @@
+#ifndef STAN_MATH_PRIM_FUN_MULTIPLY_ADD_HPP
+#define STAN_MATH_PRIM_FUN_MULTIPLY_ADD_HPP
+
+#include <stan/math/prim/meta.hpp>
+#include <stan/math/prim/err.hpp>
+#include <stan/math/prim/fun/Eigen.hpp>
+
+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 <typename Mat1, typename Mat2, typename Mat3,
+          require_all_eigen_t<Mat1, Mat2, Mat3>* = nullptr,
+          require_all_st_arithmetic<Mat1, Mat2, Mat3>* = nullptr>
+inline auto multiply_add(Mat1&& A, Mat2&& B, Mat3&& C) {
+  check_multiplicable("multiply_add", "A", A, "B", B);
+  check_size_match("multiply_add", "Rows of A", A.rows(), "rows of C",
+                   C.rows());
+  check_size_match("multiply_add", "Columns of B", B.cols(), "columns of C",
+                   C.cols());
+  return make_holder(
+      [](auto&& A_, auto&& B_, auto&& C_) {
+        return std::forward<decltype(A_)>(A_) * std::forward<decltype(B_)>(B_)
+               + std::forward<decltype(C_)>(C_);
+      },
+      std::forward<Mat1>(A), std::forward<Mat2>(B), std::forward<Mat3>(C));
+}
+
+}  // namespace math
+}  // namespace stan
+
+#endif

stan/math/prim/fun/multiply_sub.hpp

@@ -0,0 +1,48 @@
+#ifndef STAN_MATH_PRIM_FUN_MULTIPLY_SUB_HPP
+#define STAN_MATH_PRIM_FUN_MULTIPLY_SUB_HPP
+
+#include <stan/math/prim/meta.hpp>
+#include <stan/math/prim/err.hpp>
+#include <stan/math/prim/fun/Eigen.hpp>
+
+namespace stan {
+namespace math {
+
+/**
+ * Return a matrix minus the product of two matrices, `C - A * B`. The
+ * argument order follows `multiply_add(A, B, C) = A * B + C`.
+ *
+ * @tparam Mat1 type of the first matrix or expression
+ * @tparam Mat2 type of the second matrix or expression
+ * @tparam Mat3 type of the matrix or expression the product is subtracted
+ * from
+ * @param A first matrix
+ * @param B second matrix
+ * @param C matrix the product is subtracted from
+ * @return `C - A * B`
+ * @throw std::invalid_argument if the number of columns of `A` does not
+ * match the number of rows of `B`, or `C` does not have the dimensions of
+ * `A * B`
+ */
+template <typename Mat1, typename Mat2, typename Mat3,
+          require_all_eigen_t<Mat1, Mat2, Mat3>* = nullptr,
+          require_all_st_arithmetic<Mat1, Mat2, Mat3>* = nullptr>
+inline auto multiply_sub(Mat1&& A, Mat2&& B, Mat3&& C) {
+  check_multiplicable("multiply_sub", "A", A, "B", B);
+  check_size_match("multiply_sub", "Rows of A", A.rows(), "rows of C",
+                   C.rows());
+  check_size_match("multiply_sub", "Columns of B", B.cols(), "columns of C",
+                   C.cols());
+  return make_holder(
+      [](auto&& A_, auto&& B_, auto&& C_) {
+        return std::forward<decltype(C_)>(C_)
+               - std::forward<decltype(A_)>(A_)
+                     * std::forward<decltype(B_)>(B_);
+      },
+      std::forward<Mat1>(A), std::forward<Mat2>(B), std::forward<Mat3>(C));
+}
+
+}  // namespace math
+}  // namespace stan
+
+#endif

stan/math/prim/fun/scale_multiply.hpp

@@ -0,0 +1,42 @@
+#ifndef STAN_MATH_PRIM_FUN_SCALE_MULTIPLY_HPP
+#define STAN_MATH_PRIM_FUN_SCALE_MULTIPLY_HPP
+
+#include <stan/math/prim/meta.hpp>
+#include <stan/math/prim/err.hpp>
+#include <stan/math/prim/fun/Eigen.hpp>
+
+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 <typename Scal, typename Mat1, typename Mat2,
+          require_stan_scalar_t<Scal>* = nullptr,
+          require_all_eigen_t<Mat1, Mat2>* = nullptr,
+          require_all_st_arithmetic<Scal, Mat1, Mat2>* = nullptr>
+inline auto scale_multiply(Scal&& s, Mat1&& A, Mat2&& B) {
+  check_multiplicable("scale_multiply", "A", A, "B", B);
+  return make_holder(
+      [s](auto&& A_, auto&& B_) {
+        return s
+               * (std::forward<decltype(A_)>(A_)
+                  * std::forward<decltype(B_)>(B_));
+      },
+      std::forward<Mat1>(A), std::forward<Mat2>(B));
+}
+
+}  // namespace math
+}  // namespace stan
+
+#endif

stan/math/prim/meta/is_eigen.hpp

@@ -261,7 +261,11 @@ using require_all_not_eigen_st = require_all_not_t<
  */
 template <typename T>
 struct is_eigen_array
-    : bool_constant<is_base_pointer_convertible<Eigen::ArrayBase, T>::value> {};
+    : bool_constant<is_base_pointer_convertible<Eigen::ArrayBase,
+                                                std::decay_t<T>>::value> {};
+
+template <typename T>
+inline constexpr bool is_eigen_array_v = is_eigen_array<T>::value;
 
 /*! \ingroup require_eigens_types */
 /*! \defgroup eigen_array_types eigen_array  */

stan/math/rev/core/Eigen_NumTraits.hpp

@@ -5,6 +5,7 @@
 #include <stan/math/prim/core.hpp>
 #include <stan/math/prim/meta/lazy_select_evaluator.hpp>
 #include <stan/math/rev/meta.hpp>
+#include <stan/math/rev/meta/eigen_product_evaluator_hooks.hpp>
 #include <stan/math/rev/core/var.hpp>
 #include <stan/math/rev/core/gevv_vvv_vari.hpp>
 #include <stan/math/rev/core/std_numeric_limits.hpp>
@@ -504,4 +505,5 @@ struct general_matrix_matrix_product<Index, stan::math::var, LhsStorageOrder,
 };
 }  // namespace internal
 }  // namespace Eigen
+
 #endif

stan/math/rev/fun.hpp

@@ -126,8 +126,10 @@
 #include <stan/math/rev/fun/modified_bessel_first_kind.hpp>
 #include <stan/math/rev/fun/modified_bessel_second_kind.hpp>
 #include <stan/math/rev/fun/multiply.hpp>
+#include <stan/math/rev/fun/multiply_add.hpp>
 #include <stan/math/rev/fun/multiply_log.hpp>
 #include <stan/math/rev/fun/multiply_lower_tri_self_transpose.hpp>
+#include <stan/math/rev/fun/multiply_sub.hpp>
 #include <stan/math/rev/fun/norm.hpp>
 #include <stan/math/rev/fun/norm1.hpp>
 #include <stan/math/rev/fun/norm2.hpp>
@@ -149,6 +151,7 @@
 #include <stan/math/rev/fun/round.hpp>
 #include <stan/math/rev/fun/rows_dot_product.hpp>
 #include <stan/math/rev/fun/rows_dot_self.hpp>
+#include <stan/math/rev/fun/scale_multiply.hpp>
 #include <stan/math/rev/fun/sd.hpp>
 #include <stan/math/rev/fun/sin.hpp>
 #include <stan/math/rev/fun/singular_values.hpp>

stan/math/rev/fun/multiply_add.hpp

@@ -0,0 +1,68 @@
+#ifndef STAN_MATH_REV_FUN_MULTIPLY_ADD_HPP
+#define STAN_MATH_REV_FUN_MULTIPLY_ADD_HPP
+
+#include <stan/math/rev/meta.hpp>
+#include <stan/math/rev/core.hpp>
+#include <stan/math/rev/fun/to_arena.hpp>
+#include <stan/math/rev/fun/value_of.hpp>
+#include <stan/math/prim/err.hpp>
+#include <stan/math/prim/fun/Eigen.hpp>
+
+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 <typename Mat1, typename Mat2, typename Mat3,
+          require_all_eigen_t<Mat1, Mat2, Mat3>* = nullptr,
+          require_any_st_var<Mat1, Mat2, Mat3>* = nullptr,
+          require_all_not_complex_t<value_type_t<Mat1>, value_type_t<Mat2>,
+                                    value_type_t<Mat3>>* = nullptr>
+inline auto multiply_add(Mat1&& A, Mat2&& B, Mat3&& C) {
+  check_multiplicable("multiply_add", "A", A, "B", B);
+  check_size_match("multiply_add", "Rows of A", A.rows(), "rows of C",
+                   C.rows());
+  check_size_match("multiply_add", "Columns of B", B.cols(), "columns of C",
+                   C.cols());
+  arena_t<Mat1> arena_A(std::forward<Mat1>(A));
+  arena_t<Mat2> arena_B(std::forward<Mat2>(B));
+  arena_t<Mat3> arena_C(std::forward<Mat3>(C));
+  using return_t
+      = return_var_matrix_t<decltype(arena_A * arena_B), Mat1, Mat2, Mat3>;
+  arena_t<return_t> res
+      = value_of(arena_A) * value_of(arena_B) + value_of(arena_C);
+  reverse_pass_callback([arena_A, arena_B, arena_C, res]() mutable {
+    // Eigen's matrix product cannot read the adjoint view directly
+    const auto res_adj = res.adj().eval();
+    if constexpr (is_autodiff_v<Mat1>) {
+      arena_A.adj() += res_adj * value_of(arena_B).transpose();
+    }
+    if constexpr (is_autodiff_v<Mat2>) {
+      arena_B.adj() += value_of(arena_A).transpose() * res_adj;
+    }
+    if constexpr (is_autodiff_v<Mat3>) {
+      arena_C.adj() += res_adj;
+    }
+  });
+  return res;
+}
+
+}  // namespace math
+}  // namespace stan
+
+#endif

stan/math/rev/fun/multiply_sub.hpp

@@ -0,0 +1,69 @@
+#ifndef STAN_MATH_REV_FUN_MULTIPLY_SUB_HPP
+#define STAN_MATH_REV_FUN_MULTIPLY_SUB_HPP
+
+#include <stan/math/rev/meta.hpp>
+#include <stan/math/rev/core.hpp>
+#include <stan/math/rev/fun/to_arena.hpp>
+#include <stan/math/rev/fun/value_of.hpp>
+#include <stan/math/prim/err.hpp>
+#include <stan/math/prim/fun/Eigen.hpp>
+
+namespace stan {
+namespace math {
+
+/**
+ * Return a matrix minus the product of two matrices, `C - A * B`, where at
+ * least one argument has `var` scalars. The whole expression is one node on
+ * the reverse-mode stack. The argument order follows
+ * `multiply_add(A, B, C) = A * B + C`.
+ *
+ * @tparam Mat1 type of the first matrix or expression
+ * @tparam Mat2 type of the second matrix or expression
+ * @tparam Mat3 type of the matrix or expression the product is subtracted
+ * from
+ * @param A first matrix
+ * @param B second matrix
+ * @param C matrix the product is subtracted from
+ * @return `C - A * B`
+ * @throw std::invalid_argument if the number of columns of `A` does not
+ * match the number of rows of `B`, or `C` does not have the dimensions of
+ * `A * B`
+ */
+template <typename Mat1, typename Mat2, typename Mat3,
+          require_all_eigen_t<Mat1, Mat2, Mat3>* = nullptr,
+          require_any_st_var<Mat1, Mat2, Mat3>* = nullptr,
+          require_all_not_complex_t<value_type_t<Mat1>, value_type_t<Mat2>,
+                                    value_type_t<Mat3>>* = nullptr>
+inline auto multiply_sub(Mat1&& A, Mat2&& B, Mat3&& C) {
+  check_multiplicable("multiply_sub", "A", A, "B", B);
+  check_size_match("multiply_sub", "Rows of A", A.rows(), "rows of C",
+                   C.rows());
+  check_size_match("multiply_sub", "Columns of B", B.cols(), "columns of C",
+                   C.cols());
+  arena_t<Mat1> arena_A(std::forward<Mat1>(A));
+  arena_t<Mat2> arena_B(std::forward<Mat2>(B));
+  arena_t<Mat3> arena_C(std::forward<Mat3>(C));
+  using return_t
+      = return_var_matrix_t<decltype(arena_A * arena_B), Mat1, Mat2, Mat3>;
+  arena_t<return_t> res
+      = value_of(arena_C) - value_of(arena_A) * value_of(arena_B);
+  reverse_pass_callback([arena_A, arena_B, arena_C, res]() mutable {
+    // Eigen's matrix product cannot read the adjoint view directly
+    const auto res_adj = res.adj().eval();
+    if constexpr (is_autodiff_v<Mat1>) {
+      arena_A.adj() -= res_adj * value_of(arena_B).transpose();
+    }
+    if constexpr (is_autodiff_v<Mat2>) {
+      arena_B.adj() -= value_of(arena_A).transpose() * res_adj;
+    }
+    if constexpr (is_autodiff_v<Mat3>) {
+      arena_C.adj() += res_adj;
+    }
+  });
+  return res;
+}
+
+}  // namespace math
+}  // namespace stan
+
+#endif