Reverse mode autodiff for diag_post_multiply

stan/math/prim/fun/diag_post_multiply.hpp

@@ -20,7 +20,8 @@ namespace math {
  * vector or row_vector.
  */
 template <typename T1, typename T2, require_eigen_t<T1>* = nullptr,
-          require_eigen_vector_t<T2>* = nullptr>
+          require_eigen_vector_t<T2>* = nullptr,
+          require_all_not_st_var<T1, T2>* = nullptr>
 auto diag_post_multiply(const T1& m1, const T2& m2) {
   check_size_match("diag_post_multiply", "m2.size()", m2.size(), "m1.cols()",
                    m1.cols());

stan/math/prim/fun/diag_pre_multiply.hpp

@@ -20,7 +20,7 @@ namespace math {
  * vector or row_vector and a matrix.
  */
 template <typename T1, typename T2, require_eigen_vector_t<T1>* = nullptr,
-          require_eigen_matrix_dynamic_t<T2>* = nullptr,
+          require_eigen_t<T2>* = nullptr,
           require_all_not_st_var<T1, T2>* = nullptr>
 auto diag_pre_multiply(const T1& m1, const T2& m2) {
   check_size_match("diag_pre_multiply", "m1.size()", m1.size(), "m2.rows()",

stan/math/rev/fun.hpp

@@ -43,6 +43,7 @@
 #include <stan/math/rev/fun/cov_matrix_constrain_lkj.hpp>
 #include <stan/math/rev/fun/determinant.hpp>
 #include <stan/math/rev/fun/diag_pre_multiply.hpp>
+#include <stan/math/rev/fun/diag_post_multiply.hpp>
 #include <stan/math/rev/fun/digamma.hpp>
 #include <stan/math/rev/fun/dims.hpp>
 #include <stan/math/rev/fun/divide.hpp>

stan/math/rev/fun/diag_post_multiply.hpp

@@ -0,0 +1,63 @@
+#ifndef STAN_MATH_REV_FUN_DIAG_POST_MULTIPLY_HPP
+#define STAN_MATH_REV_FUN_DIAG_POST_MULTIPLY_HPP
+
+#include <stan/math/prim/fun/Eigen.hpp>
+#include <stan/math/prim/err.hpp>
+#include <stan/math/rev/core.hpp>
+
+namespace stan {
+namespace math {
+
+/**
+ * Return the product of the matrix and a diagonal matrix formed from the vector
+ * or row_vector.
+ *
+ * @tparam T1 type of the matrix
+ * @tparam T2 type of the vector/row_vector
+ * @param m1 input matrix
+ * @param m2 input vector/row_vector
+ *
+ * @return product of the matrix and the diagonal matrix formed from the
+ * vector or row_vector.
+ */
+template <typename T1, typename T2, require_matrix_t<T1>* = nullptr,
+          require_vector_t<T2>* = nullptr,
+          require_any_st_var<T1, T2>* = nullptr>
+auto diag_post_multiply(const T1& m1, const T2& m2) {
+  check_size_match("diag_post_multiply", "m2.size()", m2.size(), "m1.cols()",
+                   m1.cols());
+  using inner_ret_type = decltype(value_of(m1) * value_of(m2).asDiagonal());
+  using ret_type = return_var_matrix_t<inner_ret_type, T1, T2>;
+
+  if (!is_constant<T1>::value && !is_constant<T2>::value) {
+    arena_t<promote_scalar_t<var, T1>> arena_m1 = m1;
+    arena_t<promote_scalar_t<var, T2>> arena_m2 = m2;
+    arena_t<ret_type> ret(arena_m1.val() * arena_m2.val().asDiagonal());
+    reverse_pass_callback([ret, arena_m1, arena_m2]() mutable {
+      arena_m2.adj() += arena_m1.val().cwiseProduct(ret.adj()).colwise().sum();
+      arena_m1.adj() += ret.adj() * arena_m2.val().asDiagonal();
+    });
+    return ret_type(ret);
+  } else if (!is_constant<T1>::value) {
+    arena_t<promote_scalar_t<var, T1>> arena_m1 = m1;
+    arena_t<promote_scalar_t<double, T2>> arena_m2 = value_of(m2);
+    arena_t<ret_type> ret(arena_m1.val() * arena_m2.asDiagonal());
+    reverse_pass_callback([ret, arena_m1, arena_m2]() mutable {
+      arena_m1.adj() += ret.adj() * arena_m2.val().asDiagonal();
+    });
+    return ret_type(ret);
+  } else if (!is_constant<T2>::value) {
+    arena_t<promote_scalar_t<double, T1>> arena_m1 = value_of(m1);
+    arena_t<promote_scalar_t<var, T2>> arena_m2 = m2;
+    arena_t<ret_type> ret(arena_m1 * arena_m2.val().asDiagonal());
+    reverse_pass_callback([ret, arena_m1, arena_m2]() mutable {
+      arena_m2.adj() += arena_m1.val().cwiseProduct(ret.adj()).colwise().sum();
+    });
+    return ret_type(ret);
+  }
+}
+
+}  // namespace math
+}  // namespace stan
+
+#endif

test/unit/math/mix/fun/diag_post_multiply_test.cpp

@@ -1,23 +1,19 @@
 #include <test/unit/math/test_ad.hpp>
 
 void expect_diag_post_multiply(const Eigen::MatrixXd& a,
-                               const Eigen::VectorXd& v,
-                               const stan::test::ad_tolerances& tols) {
+                               const Eigen::VectorXd& v) {
   auto f = [](const auto& x, const auto& y) {
     return stan::math::diag_post_multiply(x, y);
   };
-  stan::test::expect_ad(tols, f, a, v);
+  stan::test::expect_ad(f, a, v);
+  stan::test::expect_ad_matvar(f, a, v);
   Eigen::RowVectorXd rv(v);
-  stan::test::expect_ad(tols, f, a, rv);
-}
-void expect_diag_post_multiply(const Eigen::MatrixXd& a,
-                               const Eigen::VectorXd& v) {
-  stan::test::ad_tolerances tols;
-  expect_diag_post_multiply(a, v, tols);
+  stan::test::expect_ad(f, a, rv);
+  stan::test::expect_ad_matvar(f, a, rv);
 }
-
 TEST(MathMixMatFun, diagPostMultiply) {
   using stan::test::relative_tolerance;
+
   // 0 x 0
   Eigen::MatrixXd a00(0, 0);
   Eigen::VectorXd u0(0);
@@ -50,28 +46,39 @@ TEST(MathMixMatFun, diagPostMultiply) {
   u3c << 1, 2, 3;
   expect_diag_post_multiply(a33c, u3c);
 
-  stan::test::ad_tolerances tols;
-  tols.hessian_hessian_ = relative_tolerance(1e-4, 2e-2);
-  tols.hessian_fvar_hessian_ = relative_tolerance(1e-4, 2e-2);
-
   Eigen::MatrixXd a33(3, 3);
   a33 << 1, 10, 100, 1000, 2, -4, 8, -16, 32;
   Eigen::VectorXd u3(3);
   u3 << -1.7, 111.2, -29.3;
-  expect_diag_post_multiply(a33, u3, tols);
+  expect_diag_post_multiply(a33, u3);
 
   Eigen::MatrixXd a33d(3, 3);
   a33d << 1, 0, 0, 0, 2, 0, 0, 0, 3;
   Eigen::VectorXd u3d(3);
   u3d << 1, 2, 3;
-  expect_diag_post_multiply(a33d, u3d, tols);
+  expect_diag_post_multiply(a33d, u3d);
 
   // error: mismatched sizes
   expect_diag_post_multiply(a33, u2);
   expect_diag_post_multiply(a22, u3);
 
-  // error: non-square
+  // non-square
   Eigen::MatrixXd b23(2, 3);
   b23 << 1, 2, 3, 4, 5, 6;
+  expect_diag_post_multiply(b23, u3);
+
+  Eigen::MatrixXd b32(3, 2);
+  b32 << 1, 2, 3, 4, 5, 6;
+  expect_diag_post_multiply(b32, u2);
+
+  Eigen::MatrixXd b13(1, 3);
+  b13 << 1, 2, 3;
+  expect_diag_post_multiply(b13, u3);
+
+  Eigen::MatrixXd b31(3, 1);
+  b31 << 1, 2, 3;
+  expect_diag_post_multiply(b31, u1);
+
+  // non-square error: mismatched sizes
   expect_diag_post_multiply(b23, u2);
 }

test/unit/math/opencl/rev/diag_post_multiply_test.cpp

@@ -37,8 +37,8 @@ TEST(OpenCL_diag_post_multiply, zero) {
 TEST(OpenCL_diag_post_multiply, prim_rev_values_large) {
   int N = 71;
 
-  Eigen::RowVectorXd a = Eigen::VectorXd::Random(N);
-  Eigen::RowVectorXd b = Eigen::VectorXd::Random(N);
+  Eigen::MatrixXd a = Eigen::MatrixXd::Random(N, N);
+  Eigen::RowVectorXd b = Eigen::RowVectorXd::Random(N);
   stan::math::test::compare_cpu_opencl_prim_rev(diag_post_multiply_functor, a,
                                                 b);
 }

test/unit/math/prim/fun/diag_post_multiply_test.cpp

@@ -43,3 +43,12 @@ TEST(MathMatrixPrimMat, diagPostMultiplyException) {
   v << 1, 2, 3;
   EXPECT_THROW(diag_post_multiply(m, v), std::invalid_argument);
 }
+
+TEST(MathMatrixPrimMat, diagPostMultiplyZero) {
+  Eigen::VectorXd in1(0);
+  Eigen::MatrixXd in2(0, 0);
+  Eigen::MatrixXd in3;
+  Eigen::MatrixXd in4(0, 0);
+  in3 = diag_post_multiply(in2, in1);
+  EXPECT_EQ(in4, in3);
+}