Avoid need to modify adj by re-working algebraically

Author: WardBrian

stan/math/rev/constraint/simplex_constrain.hpp

@@ -87,12 +87,15 @@ inline auto simplex_constrain(const T& y, scalar_type_t<T>& lp) {
   reverse_pass_callback([arena_y, arena_x, lp]() mutable {
     auto&& res_val = arena_x.val();
 
-    // backprop for log jacobian contribution to log density
-    arena_x.adj().array() += lp.adj() / res_val.array();
+    // backprop for log jacobian contribution to log density is equivalent to
+    // arena_x.adj().array() += lp.adj() / res_val.array();
+    // but is folded into the following to avoid needing to modify the adjoints
+    // in-place
 
     // backprop for softmax
-    Eigen::VectorXd 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.size() * lp.adj())
+          + (res_val.cwiseProduct(arena_x.adj()).array() + lp.adj()).matrix();
 
     // backprop for sum_to_zero_constrain
     internal::sum_to_zero_vector_backprop(arena_y.adj(), x_pre_softmax_adj);

stan/math/rev/constraint/stochastic_column_constrain.hpp

@@ -95,18 +95,18 @@ inline plain_type_t<T> stochastic_column_constrain(const T& y,
     const auto M = arena_y.cols();
 
     auto&& x_val = arena_x.val_op();
-
-    // backprop for log jacobian contribution to log density
-    arena_x.adj().array() += lp.adj() / x_val.array();
-
     auto&& x_adj = arena_x.adj_op();
 
+    const auto x_val_rows = x_val.rows();
+
     Eigen::VectorXd x_pre_softmax_adj(x_val.rows());
     for (Eigen::Index i = 0; i < M; ++i) {
       // backprop for softmax
       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));
+          = -x_val.col(i)
+                * (x_adj.col(i).dot(x_val.col(i)) + lp.adj() * x_val_rows)
+            + (x_val.col(i).cwiseProduct(x_adj.col(i)).array() + lp.adj())
+                  .matrix();
 
       // backprop for sum_to_zero_constrain
       internal::sum_to_zero_vector_backprop(arena_y.col(i).adj(),

stan/math/rev/constraint/stochastic_row_constrain.hpp

@@ -93,17 +93,18 @@ inline plain_type_t<T> stochastic_row_constrain(const T& y,
     const auto N = arena_y.rows();
 
     auto&& x_val = arena_x.val_op();
-    // backprop for log jacobian contribution to log density
-    arena_x.adj().array() += lp.adj() / x_val.array();
-
     auto&& x_adj = arena_x.adj_op();
 
+    const auto x_val_cols = x_val.cols();
+
     Eigen::VectorXd x_pre_softmax_adj(x_val.cols());
     for (Eigen::Index i = 0; i < N; ++i) {
       // backprop for softmax
       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));
+          = -x_val.row(i)
+                * (x_adj.row(i).dot(x_val.row(i)) + lp.adj() * x_val_cols)
+            + (x_val.row(i).cwiseProduct(x_adj.row(i)).array() + lp.adj())
+                  .matrix();
 
       // backprop for sum_to_zero_constrain
       internal::sum_to_zero_vector_backprop(arena_y.row(i).adj(),