Changes per review comments

Author: WardBrian

stan/math/prim/constraint/simplex_constrain.hpp

@@ -65,9 +65,7 @@ inline plain_type_t<Vec> simplex_constrain(const Vec& y) {
   max_val = fmax(max_val_old, z.coeff(0));
   d = d * exp(max_val_old - max_val) + exp(z.coeff(0) - max_val);
 
-  for (int i = 0; i <= N; ++i) {
-    z.coeffRef(i) = exp(z.coeff(i) - max_val) / d;
-  }
+  z.array() = (z.array() - max_val).exp() / d;
 
   return z;
 }
@@ -128,9 +126,7 @@ inline plain_type_t<Vec> simplex_constrain(const Vec& y, Lp& lp) {
   max_val = fmax(max_val_old, z.coeff(0));
   d = d * exp(max_val_old - max_val) + exp(z.coeff(0) - max_val);
 
-  for (int i = 0; i <= N; ++i) {
-    z.coeffRef(i) = exp(z.coeff(i) - max_val) / d;
-  }
+  z.array() = (z.array() - max_val).exp() / d;
 
   // equivalent to z.log().sum() + 0.5 * log(N + 1)
   lp += -(N + 1) * (max_val + log(d)) + 0.5 * log(N + 1);

stan/math/rev/constraint/stochastic_column_constrain.hpp

@@ -41,12 +41,13 @@ inline plain_type_t<T> stochastic_column_constrain(const T& y) {
   reverse_pass_callback([arena_y, arena_x]() mutable {
     const auto M = arena_y.cols();
 
-    const auto& x_val = to_ref(arena_x.val_op());
-    const auto& x_adj = to_ref(arena_x.adj_op());
+    auto&& x_val = arena_x.val_op();
+    auto&& x_adj = arena_x.adj_op();
 
+    Eigen::VectorXd x_pre_softmax_adj(x_val.rows());
     for (Eigen::Index i = 0; i < M; ++i) {
       // backprop for softmax
-      Eigen::VectorXd x_pre_softmax_adj
+      x_pre_softmax_adj.noalias()
           = -x_val.col(i) * x_adj.col(i).dot(x_val.col(i))
             + x_val.col(i).cwiseProduct(x_adj.col(i));
 
@@ -93,16 +94,17 @@ inline plain_type_t<T> stochastic_column_constrain(const T& y,
   reverse_pass_callback([arena_y, arena_x, lp]() mutable {
     const auto M = arena_y.cols();
 
-    const auto& x_val = to_ref(arena_x.val_op());
+    auto&& x_val = arena_x.val_op();
 
     // backprop for log jacobian contribution to log density
     arena_x.adj().array() += lp.adj() / x_val.array();
 
-    const auto& x_adj = to_ref(arena_x.adj_op());
+    auto&& x_adj = arena_x.adj_op();
 
+    Eigen::VectorXd x_pre_softmax_adj(x_val.rows());
     for (Eigen::Index i = 0; i < M; ++i) {
       // backprop for softmax
-      Eigen::VectorXd x_pre_softmax_adj
+      x_pre_softmax_adj.noalias()
           = -x_val.col(i) * x_adj.col(i).dot(x_val.col(i))
             + x_val.col(i).cwiseProduct(x_adj.col(i));
 

stan/math/rev/constraint/stochastic_row_constrain.hpp

@@ -39,12 +39,13 @@ inline auto stochastic_row_constrain(const T& y) {
   reverse_pass_callback([arena_y, arena_x]() mutable {
     const auto N = arena_y.rows();
 
-    const auto& x_val = to_ref(arena_x.val_op());
-    const auto& x_adj = to_ref(arena_x.adj_op());
+    auto&& x_val = arena_x.val_op();
+    auto&& x_adj = arena_x.adj_op();
 
+    Eigen::VectorXd x_pre_softmax_adj(x_val.cols());
     for (Eigen::Index i = 0; i < N; ++i) {
       // backprop for softmax
-      Eigen::VectorXd x_pre_softmax_adj
+      x_pre_softmax_adj.noalias()
           = -x_val.row(i) * x_adj.row(i).dot(x_val.row(i))
             + x_val.row(i).cwiseProduct(x_adj.row(i));
 
@@ -91,16 +92,16 @@ inline plain_type_t<T> stochastic_row_constrain(const T& y,
   reverse_pass_callback([arena_y, arena_x, lp]() mutable {
     const auto N = arena_y.rows();
 
-    const auto& x_val = to_ref(arena_x.val_op());
-
+    auto&& x_val = arena_x.val_op();
     // backprop for log jacobian contribution to log density
     arena_x.adj().array() += lp.adj() / x_val.array();
 
-    const auto& x_adj = to_ref(arena_x.adj_op());
+    auto&& x_adj = arena_x.adj_op();
 
+    Eigen::VectorXd x_pre_softmax_adj(x_val.cols());
     for (Eigen::Index i = 0; i < N; ++i) {
       // backprop for softmax
-      Eigen::VectorXd x_pre_softmax_adj
+      x_pre_softmax_adj.noalias()
           = -x_val.row(i) * x_adj.row(i).dot(x_val.row(i))
             + x_val.row(i).cwiseProduct(x_adj.row(i));
 

stan/math/rev/constraint/sum_to_zero_constrain.hpp

@@ -21,6 +21,10 @@ namespace internal {
  * The reverse pass backprop for the sum_to_zero_constrain on
  * vectors. This is separated out so it can also be called by
  * simplex_constrain.
+ *
+ * @tparam T type of the adjoint vector
+ * @param y_adj The adjoint of the free vector (size N)
+ * @param z_adj The adjoint of the zero-sum vector (size N + 1)
  */
 template <typename T>
 void sum_to_zero_vector_backprop(T&& y_adj, const Eigen::VectorXd& z_adj) {