Write out derivatives

Author: WardBrian

stan/math/rev/constraint/simplex_constrain.hpp

@@ -5,10 +5,8 @@
 #include <stan/math/rev/meta.hpp>
 #include <stan/math/rev/core/reverse_pass_callback.hpp>
 #include <stan/math/rev/core/arena_matrix.hpp>
-#include <stan/math/rev/fun/value_of.hpp>
-#include <stan/math/prim/constraint/sum_to_zero_constrain.hpp>
-#include <stan/math/prim/fun/softmax.hpp>
-#include <stan/math/prim/fun/log_softmax.hpp>
+#include <stan/math/rev/constraint/sum_to_zero_constrain.hpp>
+#include <stan/math/prim/constraint/simplex_constrain.hpp>
 #include <cmath>
 #include <tuple>
 #include <vector>
@@ -32,7 +30,30 @@ namespace math {
  */
 template <typename T, require_rev_col_vector_t<T>* = nullptr>
 inline auto simplex_constrain(const T& y) {
-  return softmax(sum_to_zero_constrain(y));
+  using ret_type = plain_type_t<T>;
+
+  const auto N = y.size();
+  arena_t<T> arena_y = y;
+
+  arena_t<ret_type> arena_x = simplex_constrain(arena_y.val());
+
+  if (unlikely(N == 0)) {
+    return ret_type(arena_x);
+  }
+
+  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());
+
+    // backprop for sum_to_zero_constrain
+    internal::sum_to_zero_vector_backprop(arena_y, x_pre_softmax_adj);
+  });
+
+  return ret_type(arena_x);
 }
 
 /**
@@ -50,16 +71,36 @@ inline auto simplex_constrain(const T& y) {
  * @return Simplex of dimensionality N + 1.
  */
 template <typename T, require_rev_col_vector_t<T>* = nullptr>
-auto simplex_constrain(const T& y, scalar_type_t<T>& lp) {
+inline auto simplex_constrain(const T& y, scalar_type_t<T>& lp) {
   using ret_type = plain_type_t<T>;
 
-  arena_t<ret_type> log_x = log_softmax(sum_to_zero_constrain(y));
+  const auto N = y.size();
+  arena_t<T> arena_y = y;
+
+  double lp_val = 0.0;
+  arena_t<ret_type> arena_x = simplex_constrain(arena_y.val(), lp_val);
+  lp += lp_val;
+
+  if (unlikely(N == 0)) {
+    return ret_type(arena_x);
+  }
+
+  reverse_pass_callback([arena_y, arena_x, lp]() mutable {
+    const auto& res_val = to_ref(arena_x.val());
+
+    // backprop for log jacobian contribution to log density
+    arena_x.adj().array() += lp.adj() / res_val.array();
 
-  const auto N = y.size() + 1;
+    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());
 
-  lp += sum(log_x) + 0.5 * log(N);
+    // backprop for sum_to_zero_constrain
+    internal::sum_to_zero_vector_backprop(arena_y, x_pre_softmax_adj);
+  });
 
-  return ret_type(exp(log_x));
+  return ret_type(arena_x);
 }
 
 }  // namespace math

stan/math/rev/constraint/sum_to_zero_constrain.hpp

@@ -15,6 +15,35 @@
 namespace stan {
 namespace math {
 
+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.
+ */
+template <typename T>
+void sum_to_zero_vector_backprop(T&& arena_y, Eigen::VectorXd z_adj) {
+  const auto N = arena_y.size();
+
+  double sum_u_adj = 0;
+  for (int i = 0; i < N; ++i) {
+    double n = static_cast<double>(i + 1);
+
+    // adjoint of the reverse cumulative sum computed in the forward mode
+    sum_u_adj += z_adj.coeff(i);
+
+    // adjoint of the offset subtraction
+    double v_adj = -z_adj.coeff(i + 1) * n;
+
+    double w_adj = v_adj + sum_u_adj;
+
+    arena_y.adj().coeffRef(i) += w_adj / sqrt(n * (n + 1));
+  }
+}
+
+}  // namespace internal
+
 /**
  * Return a vector with sum zero corresponding to the specified
  * free vector.
@@ -49,22 +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 {
-    const auto N = arena_y.size();
-
-    double sum_u_adj = 0;
-    for (int i = 0; i < N; ++i) {
-      double n = static_cast<double>(i + 1);
-
-      // adjoint of the reverse cumulative sum computed in the forward mode
-      sum_u_adj += arena_z.adj().coeff(i);
-
-      // adjoint of the offset subtraction
-      double v_adj = -arena_z.adj().coeff(i + 1) * n;
-
-      double w_adj = v_adj + sum_u_adj;
-
-      arena_y.adj().coeffRef(i) += w_adj / sqrt(n * (n + 1));
-    }
+    internal::sum_to_zero_vector_backprop(arena_y, arena_z.adj());
   });
 
   return arena_z;