Clean up simplex rev

Author: WardBrian

stan/math/rev/constraint/simplex_constrain.hpp

@@ -44,13 +44,12 @@ inline auto simplex_constrain(const T& y) {
   reverse_pass_callback([arena_y, arena_x]() mutable {
     const auto& res_val = to_ref(arena_x.val());
 
-    Eigen::VectorXd x_pre_softmax_adj = Eigen::VectorXd::Zero(res_val.size());
     // backprop for softmax
-    x_pre_softmax_adj += -res_val * arena_x.adj().dot(res_val)
-                         + res_val.cwiseProduct(arena_x.adj());
+    Eigen::VectorXd x_pre_softmax_adj = -res_val * arena_x.adj().dot(res_val)
+                                        + res_val.cwiseProduct(arena_x.adj());
 
     // backprop for sum_to_zero_constrain
-    internal::sum_to_zero_vector_backprop(arena_y, x_pre_softmax_adj);
+    internal::sum_to_zero_vector_backprop(arena_y.adj(), x_pre_softmax_adj);
   });
 
   return ret_type(arena_x);
@@ -91,13 +90,12 @@ inline auto simplex_constrain(const T& y, scalar_type_t<T>& lp) {
     // backprop for log jacobian contribution to log density
     arena_x.adj().array() += lp.adj() / res_val.array();
 
-    Eigen::VectorXd x_pre_softmax_adj = Eigen::VectorXd::Zero(res_val.size());
     // backprop for softmax
-    x_pre_softmax_adj += -res_val * arena_x.adj().dot(res_val)
-                         + res_val.cwiseProduct(arena_x.adj());
+    Eigen::VectorXd x_pre_softmax_adj = -res_val * arena_x.adj().dot(res_val)
+                                        + res_val.cwiseProduct(arena_x.adj());
 
     // backprop for sum_to_zero_constrain
-    internal::sum_to_zero_vector_backprop(arena_y, x_pre_softmax_adj);
+    internal::sum_to_zero_vector_backprop(arena_y.adj(), x_pre_softmax_adj);
   });
 
   return ret_type(arena_x);

stan/math/rev/constraint/sum_to_zero_constrain.hpp

@@ -23,8 +23,8 @@ namespace internal {
  * simplex_constrain.
  */
 template <typename T>
-void sum_to_zero_vector_backprop(T&& arena_y, Eigen::VectorXd z_adj) {
-  const auto N = arena_y.size();
+void sum_to_zero_vector_backprop(T&& y_adj, const Eigen::VectorXd& z_adj) {
+  const auto N = y_adj.size();
 
   double sum_u_adj = 0;
   for (int i = 0; i < N; ++i) {
@@ -38,7 +38,7 @@ void sum_to_zero_vector_backprop(T&& arena_y, Eigen::VectorXd z_adj) {
 
     double w_adj = v_adj + sum_u_adj;
 
-    arena_y.adj().coeffRef(i) += w_adj / sqrt(n * (n + 1));
+    y_adj.coeffRef(i) += w_adj / sqrt(n * (n + 1));
   }
 }
 
@@ -78,7 +78,7 @@ inline auto sum_to_zero_constrain(T&& y) {
   arena_t<ret_type> arena_z = sum_to_zero_constrain(arena_y.val());
 
   reverse_pass_callback([arena_y, arena_z]() mutable {
-    internal::sum_to_zero_vector_backprop(arena_y, arena_z.adj());
+    internal::sum_to_zero_vector_backprop(arena_y.adj(), arena_z.adj());
   });
 
   return arena_z;