Skeleton of rev

Author: WardBrian

stan/math/prim/constraint/sum_to_zero_constrain.hpp

@@ -83,7 +83,7 @@ inline plain_type_t<Mat> sum_to_zero_constrain(const Mat& x) {
   }
   auto&& x_ref = to_ref(x);
 
-  Eigen::VectorXd beta = Eigen::VectorXd::Zero(N);
+  Eigen::Matrix<value_type_t<Mat>, -1, 1> beta = Eigen::VectorXd::Zero(N);
 
   for (int j = M - 1; j >= 0; --j) {
     value_type_t<Mat> ax_previous(0);

stan/math/rev/constraint/sum_to_zero_constrain.hpp

@@ -69,6 +69,34 @@ inline auto sum_to_zero_constrain(T&& y) {
   return arena_z;
 }
 
+/**
+ * Return a matrix that sums to zero over both the rows
+ * and columns corresponding to the free matrix x.
+ *
+ * This is a linear transform, with no Jacobian.
+ *
+ * @tparam Mat type of the matrix
+ * @param x Free matrix input of dimensionality (N - 1, M - 1).
+ * @return Zero-sum matrix of dimensionality (N, M).
+ */
+template <typename T, require_rev_matrix_t<T>* = nullptr,
+          require_not_t<is_rev_vector<T>>* = nullptr>
+inline auto sum_to_zero_constrain(T&& x) {
+  using ret_type = plain_type_t<T>;
+  if (unlikely(x.size() == 0)) {
+    return arena_t<ret_type>(Eigen::MatrixXd{{0}});
+  }
+  auto arena_x = to_arena(std::forward<T>(x));
+  arena_t<ret_type> arena_z = sum_to_zero_constrain(arena_x.val());
+
+  reverse_pass_callback([arena_x, arena_z]() mutable {
+    const auto N = arena_x.rows();
+    const auto M = arena_x.cols();
+  });
+
+  return arena_z;
+}
+
 /**
  * Return a vector with sum zero corresponding to the specified
  * free vector.
@@ -89,12 +117,12 @@ inline auto sum_to_zero_constrain(T&& y) {
  *
  * This is a linear transform, with no Jacobian.
  *
- * @tparam Vec type of the vector
- * @param y Free vector input of dimensionality K - 1.
+ * @tparam T type of the vector or matrix
+ * @param y Free vector or matrix.
  * @param lp unused
- * @return Zero-sum vector of dimensionality K.
+ * @return Zero-sum vector or matrix which is one larger in each dimension
  */
-template <typename T, typename Lp, require_rev_col_vector_t<T>* = nullptr>
+template <typename T, typename Lp, is_rev_matrix<T>* = nullptr>
 inline auto sum_to_zero_constrain(T&& y, Lp& lp) {
   return sum_to_zero_constrain(std::forward<T>(y));
 }

test/unit/math/mix/constraint/sum_to_zero_constrain_test.cpp

@@ -55,4 +55,20 @@ TEST(MathMixMatFun, sum_to_zeroTransform) {
   Eigen::VectorXd v5(5);
   v5 << 1, -3, 2, 0, -1;
   sum_to_zero_constrain_test::expect_sum_to_zero_transform(v5);
+
+  Eigen::MatrixXd m0_0(0, 0);
+  sum_to_zero_constrain_test::expect_sum_to_zero_transform(m0_0);
+
+  Eigen::MatrixXd m1_1(1, 1);
+  m1_1 << 1;
+  sum_to_zero_constrain_test::expect_sum_to_zero_transform(m1_1);
+
+  Eigen::MatrixXd m2_2(2, 2);
+  m2_2 << 1, 2, -3, 4;
+  sum_to_zero_constrain_test::expect_sum_to_zero_transform(m2_2);
+
+  Eigen::MatrixXd m3_4(3, 4);
+  m3_4 << 1, 2, -3, 4, 5, 6, -7, 8, 9, -10, 11, -12;
+
+  sum_to_zero_constrain_test::expect_sum_to_zero_transform(m3_4);
 }