Fix scale_matrix_exp_multiply to use to_ref and make_holder

Author: SteveBronder

stan/math/fwd/fun/multiply.hpp

@@ -17,28 +17,34 @@ template <typename Mat1, typename Mat2,
           require_all_eigen_vt<is_fvar, Mat1, Mat2>* = nullptr,
           require_vt_same<Mat1, Mat2>* = nullptr,
           require_not_eigen_row_and_col_t<Mat1, Mat2>* = nullptr>
-inline auto multiply(const Mat1& m1, const Mat2& m2) {
+inline auto multiply(Mat1&& m1, Mat2&& m2) {
   check_multiplicable("multiply", "m1", m1, "m2", m2);
-  return to_fvar(multiply(m1.val(), m2.val()),
-                 add(multiply(m1.val(), m2.d()), multiply(m1.d(), m2.val())));
+  decltype(auto) m1_ref = to_ref(std::forward<Mat1>(m1));
+  decltype(auto) m2_ref = to_ref(std::forward<Mat2>(m2));
+  return to_fvar(multiply(m1_ref.val(), m2_ref.val()),
+                 add(multiply(m1_ref.val(), m2_ref.d()), multiply(m1_ref.d(), m2_ref.val())));
 }
 
 template <typename Mat1, typename Mat2,
           require_eigen_vt<is_fvar, Mat1>* = nullptr,
           require_eigen_vt<std::is_floating_point, Mat2>* = nullptr,
           require_not_eigen_row_and_col_t<Mat1, Mat2>* = nullptr>
-inline auto multiply(const Mat1& m1, const Mat2& m2) {
+inline auto multiply(Mat1&& m1, Mat2&& m2) {
   check_multiplicable("multiply", "m1", m1, "m2", m2);
-  return to_fvar(multiply(m1.val(), m2), multiply(m1.d(), m2));
+  decltype(auto) m1_ref = to_ref(std::forward<Mat1>(m1));
+  decltype(auto) m2_ref = to_ref(std::forward<Mat2>(m2));
+  return to_fvar(multiply(m1_ref.val(), m2_ref), multiply(m1_ref.d(), m2_ref));
 }
 
 template <typename Mat1, typename Mat2,
           require_eigen_vt<std::is_floating_point, Mat1>* = nullptr,
           require_eigen_vt<is_fvar, Mat2>* = nullptr,
           require_not_eigen_row_and_col_t<Mat1, Mat2>* = nullptr>
-inline auto multiply(const Mat1& m1, const Mat2& m2) {
+inline auto multiply(Mat1&& m1, Mat2&& m2) {
   check_multiplicable("multiply", "m1", m1, "m2", m2);
-  return to_fvar(multiply(m1, m2.val()), multiply(m1, m2.d()));
+  decltype(auto) m1_ref = to_ref(std::forward<Mat1>(m1));
+  decltype(auto) m2_ref = to_ref(std::forward<Mat2>(m2));
+  return to_fvar(multiply(m1_ref, m2_ref.val()), multiply(m1_ref, m2_ref.d()));
 }
 
 }  // namespace math

stan/math/prim/fun/matrix_exp.hpp

@@ -21,17 +21,17 @@ namespace math {
  * @throw <code>std::invalid_argument</code> if the input matrix
  * is not square.
  */
-template <typename T, typename = require_eigen_t<T>>
-inline plain_type_t<T> matrix_exp(const T& A_in) {
-  using std::exp;
-  const auto& A = A_in.eval();
+template <typename EigenMat, typename = require_eigen_t<EigenMat>>
+inline plain_type_t<EigenMat> matrix_exp(EigenMat&& A_in) {
+  decltype(auto) A = to_ref(std::forward<EigenMat>(A_in));
+  using T = std::decay_t<EigenMat>;
   check_square("matrix_exp", "input matrix", A);
   if constexpr (T::RowsAtCompileTime == 1 && T::ColsAtCompileTime == 1) {
     plain_type_t<T> res;
     res << exp(A(0));
     return res;
   }
-  if (A_in.size() == 0) {
+  if (A.size() == 0) {
     return {};
   }
   return (A.cols() == 2

stan/math/prim/fun/matrix_exp_action_handler.hpp

@@ -4,6 +4,7 @@
 #include <stan/math/prim/meta.hpp>
 #include <stan/math/prim/fun/Eigen.hpp>
 #include <stan/math/prim/fun/ceil.hpp>
+#include <stan/math/prim/fun/to_ref.hpp>
 #include <unsupported/Eigen/MatrixFunctions>
 #include <vector>
 #include <cmath>
@@ -23,10 +24,10 @@ namespace math {
  * and t is double.
  */
 class matrix_exp_action_handler {
-  const int _p_max = 8;
-  const int _m_max = 55;
-  const double _tol = 1.1e-16;  // from the paper, double precision: 2^-53
-  const std::vector<double> _theta_m{
+  static constexpr int _p_max = 8;
+  static constexpr int _m_max = 55;
+  static constexpr double _tol = 1.1e-16;  // from the paper, double precision: 2^-53
+  static constexpr std::array<double, 100> _theta_m{
       2.22044605e-16, 2.58095680e-08, 1.38634787e-05, 3.39716884e-04,
       2.40087636e-03, 9.06565641e-03, 2.38445553e-02, 4.99122887e-02,
       8.95776020e-02, 1.44182976e-01, 2.14235807e-01, 2.99615891e-01,
@@ -64,18 +65,18 @@ class matrix_exp_action_handler {
   template <typename EigMat1, typename EigMat2,
             require_all_eigen_t<EigMat1, EigMat2>* = nullptr,
             require_all_st_same<double, EigMat1, EigMat2>* = nullptr>
-  inline Eigen::MatrixXd action(const EigMat1& mat, const EigMat2& b,
-                                const double& t = 1.0) {
-    Eigen::MatrixXd A = mat;
+  inline Eigen::MatrixXd action(EigMat1&& mat, EigMat2&& b,
+                                const double t = 1.0) {
+    decltype(auto) A = to_ref(std::forward<EigMat1>(mat));
     double mu = A.trace() / A.rows();
     A.diagonal().array() -= mu;
 
     int m, s;
-    set_approx_order(A, b, t, m, s);
+    decltype(auto) b_eval = to_ref(std::forward<EigMat2>(b));
+    set_approx_order(A, b_eval, t, m, s);
 
     double eta = exp(t * mu / s);
 
-    const auto& b_eval = b.eval();
     Eigen::MatrixXd f = b_eval;
     Eigen::MatrixXd bi = b_eval;
 
@@ -102,8 +103,11 @@ class matrix_exp_action_handler {
    *
    * @param x matrix
    */
-  double matrix_operator_inf_norm(Eigen::MatrixXd const& x) {
-    return x.cwiseAbs().rowwise().sum().maxCoeff();
+  template <typename EigenMat>
+  double matrix_operator_inf_norm(EigenMat&& x) {
+    return make_holder([](auto&& x_) {
+    return x_.cwiseAbs().rowwise().sum().maxCoeff();
+    }, std::forward<EigenMat>(x));
   }
 
   /**
@@ -125,15 +129,16 @@ class matrix_exp_action_handler {
    */
   template <typename EigMat1, require_all_eigen_t<EigMat1>* = nullptr,
             require_all_st_same<double, EigMat1>* = nullptr>
-  double mat_power_1_norm(const EigMat1& mat, int m) {
-    if ((mat.array() > 0.0).all()) {
-      Eigen::VectorXd e = Eigen::VectorXd::Constant(mat.rows(), 1.0);
+  double mat_power_1_norm(EigMat1&& mat, const int m) {
+    auto&& mat_ref = to_ref(std::forward<EigMat1>(mat));
+    if ((mat_ref.array() > 0.0).all()) {
+      Eigen::VectorXd e = Eigen::VectorXd::Constant(mat_ref.rows(), 1.0);
       for (int j = 0; j < m; ++j) {
-        e = mat.transpose() * e;
+        e = mat_ref.transpose() * e;
       }
       return e.lpNorm<Eigen::Infinity>();
     } else {
-      return mat.pow(m).cwiseAbs().colwise().sum().maxCoeff();
+      return mat_ref.pow(m).cwiseAbs().colwise().sum().maxCoeff();
     }
   }
 
@@ -156,8 +161,8 @@ class matrix_exp_action_handler {
   template <typename EigMat1, typename EigMat2,
             require_all_eigen_t<EigMat1, EigMat2>* = nullptr,
             require_all_st_same<double, EigMat1, EigMat2>* = nullptr>
-  inline void set_approx_order(const EigMat1& mat, const EigMat2& b,
-                               const double& t, int& m, int& s) {
+  inline void set_approx_order(EigMat1&& mat, EigMat2&& b,
+                               const double t, int& m, int& s) {
     if (t < _tol) {
       m = 0;
       s = 1;

stan/math/prim/fun/scale_matrix_exp_multiply.hpp

@@ -26,15 +26,15 @@ namespace math {
  */
 template <typename EigMat1, typename EigMat2,
           require_all_eigen_vt<std::is_arithmetic, EigMat1, EigMat2>* = nullptr>
-inline Eigen::Matrix<double, Eigen::Dynamic, EigMat2::ColsAtCompileTime>
-scale_matrix_exp_multiply(const double& t, const EigMat1& A, const EigMat2& B) {
+inline Eigen::Matrix<double, Eigen::Dynamic, std::decay_t<EigMat2>::ColsAtCompileTime>
+scale_matrix_exp_multiply(const double t, EigMat1&& A, EigMat2&& B) {
   check_square("scale_matrix_exp_multiply", "input matrix", A);
   check_multiplicable("scale_matrix_exp_multiply", "A", A, "B", B);
   if (A.size() == 0) {
     return {0, B.cols()};
   }
 
-  return matrix_exp_action_handler().action(A, B, t);
+  return matrix_exp_action_handler().action(std::forward<EigMat1>(A), std::forward<EigMat2>(B), t);
 }
 
 /**
@@ -56,15 +56,15 @@ template <typename Tt, typename EigMat1, typename EigMat2,
           require_any_autodiff_scalar_t<Tt, value_type_t<EigMat1>,
                                         value_type_t<EigMat2>>* = nullptr>
 inline Eigen::Matrix<return_type_t<Tt, EigMat1, EigMat2>, Eigen::Dynamic,
-                     EigMat2::ColsAtCompileTime>
-scale_matrix_exp_multiply(const Tt& t, const EigMat1& A, const EigMat2& B) {
+                     std::decay_t<EigMat2>::ColsAtCompileTime>
+scale_matrix_exp_multiply(const Tt t, EigMat1&& A, EigMat2&& B) {
   check_square("scale_matrix_exp_multiply", "input matrix", A);
   check_multiplicable("scale_matrix_exp_multiply", "A", A, "B", B);
   if (A.size() == 0) {
     return {0, B.cols()};
   }
 
-  return multiply(matrix_exp(multiply(A, t)), B);
+  return multiply(matrix_exp(multiply(std::forward<EigMat1>(A), t)), std::forward<EigMat2>(B));
 }
 
 }  // namespace math