Support softmax for row_vector, matrix types + code simplifications

stan/math/fwd/fun/log_softmax.hpp

@@ -5,53 +5,53 @@
 #include <stan/math/fwd/core.hpp>
 #include <stan/math/fwd/meta.hpp>
 #include <stan/math/fwd/fun/softmax.hpp>
+#include <stan/math/prim/err.hpp>
 #include <stan/math/prim/fun/log_softmax.hpp>
 #include <stan/math/prim/fun/to_ref.hpp>
+#include <stan/math/prim/functor/apply_vector_unary.hpp>
 
 namespace stan {
 namespace math {
 
 /**
- * Return the log softmax of the specified vector or container of vectors.
+ * Return the log softmax of each vector in a container of `fvar` values.
  *
- * @tparam T Type of input vector or matrix.
- * @param[in] x Unconstrained input vector.
- * @return Softmax of the input.
- * @throw std::domain_error If the input vector is size 0.
+ * @tparam T `std::vector` whose scalar type is `fvar`
+ * @param x container of vectors to transform
+ * @return container of log softmax results
  */
-template <typename T, require_vector_st<is_fvar, T>* = nullptr>
+template <typename T, require_std_vector_st<is_fvar, T>* = nullptr>
 inline auto log_softmax(T&& x) {
-  return apply_vector_unary<T>::apply(std::forward<T>(x), [](auto&& alpha) {
-    using T_alpha = decltype(alpha);
-    using T_fvar = value_type_t<T_alpha>;
-    using T_fvar_inner = typename T_fvar::Scalar;
-
-    auto&& alpha_ref = to_ref(std::forward<decltype(alpha)>(alpha));
-    Eigen::Matrix<T_fvar_inner, -1, 1> alpha_t = alpha_ref.val();
-    Eigen::Matrix<T_fvar_inner, -1, 1> softmax_alpha_t = softmax(alpha_t);
-
-    Eigen::Matrix<T_fvar, -1, 1> log_softmax_alpha(alpha_ref.size());
-    log_softmax_alpha.val() = log_softmax(alpha_t);
-    log_softmax_alpha.d().setZero();
-
-    for (int m = 0; m < alpha_ref.size(); ++m) {
-      T_fvar_inner negative_alpha_m_d_times_softmax_alpha_t_m
-          = -alpha_ref.coeff(m).d_ * softmax_alpha_t(m);
-      for (int k = 0; k < alpha_ref.size(); ++k) {
-        if (m == k) {
-          log_softmax_alpha(k).d_
-              += alpha_ref.coeff(m).d_
-                 + negative_alpha_m_d_times_softmax_alpha_t_m;
-        } else {
-          log_softmax_alpha(k).d_ += negative_alpha_m_d_times_softmax_alpha_t_m;
-        }
-      }
-    }
-
-    return log_softmax_alpha;
+  return apply_vector_unary<T>::apply(std::forward<T>(x), [](auto&& v) {
+    return log_softmax(std::forward<decltype(v)>(v));
   });
 }
 
+/**
+ * Return the log softmax of the specified vector of `fvar` values.
+ *
+ * @tparam Vec Eigen vector with `fvar` scalar
+ * @param x vector to transform
+ * @return log softmax of the vector
+ * @throw std::domain_error if the input size is 0
+ */
+template <typename Vec, require_eigen_vector_vt<is_fvar, Vec>* = nullptr>
+inline auto log_softmax(Vec&& x) {
+  using vec = std::decay_t<Vec>;
+  constexpr int Rows = vec::RowsAtCompileTime;
+  constexpr int Cols = vec::ColsAtCompileTime;
+  using T = typename value_type_t<Vec>::Scalar;
+  check_nonzero_size("log_softmax", "x", x);
+  decltype(auto) x_ref = to_ref(std::forward<Vec>(x));
+  const auto s = softmax(value_of(x_ref));
+  const auto d_in = x_ref.d();
+  const auto dot_sd = s.dot(d_in);
+  Eigen::Matrix<fvar<T>, Rows, Cols> result(x_ref.size());
+  result.val() = s.array().log().matrix();
+  result.d() = (d_in.array() - dot_sd).matrix();
+  return result;
+}
+
 }  // namespace math
 }  // namespace stan
 #endif

stan/math/fwd/fun/softmax.hpp

@@ -6,47 +6,49 @@
 #include <stan/math/fwd/fun/value_of.hpp>
 #include <stan/math/prim/fun/to_ref.hpp>
 #include <stan/math/prim/fun/softmax.hpp>
+#include <stan/math/prim/functor/apply_vector_unary.hpp>
 
 namespace stan {
 namespace math {
 
-template <typename ColVec,
-          require_eigen_col_vector_vt<is_fvar, ColVec>* = nullptr>
-inline auto softmax(const ColVec& alpha) {
-  using Eigen::Dynamic;
-  using Eigen::Matrix;
-  using T = typename value_type_t<ColVec>::Scalar;
-  if (alpha.size() == 0) {
-    return Matrix<fvar<T>, Dynamic, 1>();
-  }
-  const auto& alpha_ref = to_ref(alpha);
-
-  Matrix<T, Dynamic, 1> softmax_alpha_t = softmax(value_of(alpha_ref));
-
-  Matrix<fvar<T>, Dynamic, 1> softmax_alpha(alpha.size());
-  for (int k = 0; k < alpha.size(); ++k) {
-    softmax_alpha.coeffRef(k).val_ = softmax_alpha_t.coeff(k);
-    softmax_alpha.coeffRef(k).d_ = 0;
-  }
+/**
+ * Return the softmax of each vector in a container of `fvar` values.
+ *
+ * @tparam T `std::vector` whose scalar type is `fvar`
+ * @param x container of vectors to transform
+ * @return container of softmax results
+ */
+template <typename T, require_std_vector_st<is_fvar, T>* = nullptr>
+inline auto softmax(T&& x) {
+  return apply_vector_unary<T>::apply(std::forward<T>(x), [](auto&& v) {
+    return softmax(std::forward<decltype(v)>(v));
+  });
+}
 
-  for (int m = 0; m < alpha.size(); ++m) {
-    T negative_alpha_m_d_times_softmax_alpha_t_m
-        = -alpha_ref.coeff(m).d_ * softmax_alpha_t.coeff(m);
-    for (int k = 0; k < alpha.size(); ++k) {
-      if (m == k) {
-        softmax_alpha.coeffRef(k).d_
-            += softmax_alpha_t.coeff(k)
-               * (alpha_ref.coeff(m).d_
-                  + negative_alpha_m_d_times_softmax_alpha_t_m);
-      } else {
-        softmax_alpha.coeffRef(k).d_
-            += softmax_alpha_t.coeff(k)
-               * negative_alpha_m_d_times_softmax_alpha_t_m;
-      }
-    }
+/**
+ * Return the softmax of the specified vector of `fvar` values.
+ *
+ * @tparam Vec Eigen vector with `fvar` scalar
+ * @param x vector to transform
+ * @return softmax of the vector
+ */
+template <typename Vec, require_eigen_vector_vt<is_fvar, Vec>* = nullptr>
+inline auto softmax(Vec&& x) {
+  using vec = std::decay_t<Vec>;
+  constexpr int Rows = vec::RowsAtCompileTime;
+  constexpr int Cols = vec::ColsAtCompileTime;
+  using T = typename value_type_t<vec>::Scalar;
+  if (x.size() == 0) {
+    return Eigen::Matrix<fvar<T>, Rows, Cols>();
   }
-
-  return softmax_alpha;
+  decltype(auto) x_ref = to_ref(std::forward<Vec>(x));
+  const auto s = softmax(value_of(x_ref));
+  const auto d_in = x_ref.d();
+  const auto dot_sd = s.dot(d_in);
+  Eigen::Matrix<fvar<T>, Rows, Cols> result(x_ref.size());
+  result.val() = s;
+  result.d() = (s.array() * (d_in.array() - dot_sd)).matrix();
+  return result;
 }
 
 }  // namespace math

stan/math/prim/fun/log_softmax.hpp

@@ -13,7 +13,7 @@ namespace math {
 
 /**
  * Return the natural logarithm of the softmax of the specified
- * vector.
+ * vector, or of each vector in a container.
  *
  * \f$
  * \log \mbox{softmax}(y)
@@ -23,26 +23,31 @@ namespace math {
  *
  * For the log softmax function, the entries in the Jacobian are
  * \f$
- * \frac{\partial}{\partial y_m} \mbox{softmax}(y)[k]
+ * \frac{\partial}{\partial y_m} \log\mbox{softmax}(y)[k]
  * = \left\{
  * \begin{array}{ll}
  * 1 - \mbox{softmax}(y)[m]
  * & \mbox{ if } m = k, \mbox{ and}
  * \\[6pt]
- * \mbox{softmax}(y)[m]
+ * -\mbox{softmax}(y)[m]
  * & \mbox{ if } m \neq k.
  * \end{array}
  * \right.
  * \f$
  *
- * @tparam Container type of input vector to transform
- * @param[in] x vector to transform
- * @return log unit simplex result of the softmax transform of the vector.
+ * @tparam Container type of input: an Eigen vector, `std::vector` of doubles,
+ *   or nested container whose scalar type is arithmetic
+ * @param[in] x vector or container of vectors to transform
+ * @return log softmax of the input, preserving the container structure
+ * @throw std::domain_error if any input vector is empty
  */
 template <typename Container, require_st_arithmetic<Container>* = nullptr,
-          require_container_t<Container>* = nullptr>
+          require_container_t<Container>* = nullptr,
+          require_not_t<bool_constant<
+              is_eigen<std::decay_t<Container>>::value
+              && !is_eigen_vector<std::decay_t<Container>>::value>>* = nullptr>
 inline auto log_softmax(Container&& x) {
-  check_nonzero_size("log_softmax", "v", x);
+  check_nonzero_size("log_softmax", "x", x);
   return make_holder(
       [](auto&& a) {
         return apply_vector_unary<ref_type_t<Container>>::apply(

stan/math/prim/fun/softmax.hpp

@@ -4,6 +4,7 @@
 #include <stan/math/prim/err.hpp>
 #include <stan/math/prim/fun/Eigen.hpp>
 #include <stan/math/prim/fun/to_ref.hpp>
+#include <stan/math/prim/functor/apply_vector_unary.hpp>
 #include <cmath>
 
 namespace stan {
@@ -38,20 +39,33 @@ namespace math {
  * \end{array}
  * \f$
  *
- * @tparam ColVec type of elements in the vector
+ * @tparam Vec type of the input vector
  * @param[in] v Vector to transform.
  * @return Unit simplex result of the softmax transform of the vector.
  */
-template <typename ColVec,
-          require_eigen_col_vector_vt<std::is_arithmetic, ColVec>* = nullptr>
-inline plain_type_t<ColVec> softmax(const ColVec& v) {
-  using std::exp;
+template <typename Vec,
+          require_eigen_vector_vt<std::is_arithmetic, Vec>* = nullptr>
+inline plain_type_t<Vec> softmax(Vec&& v) {
   if (v.size() == 0) {
     return v;
   }
-  const auto& v_ref = to_ref(v);
-  const auto theta = (v_ref.array() - v_ref.maxCoeff()).exp().eval();
-  return theta.array() / theta.sum();
+  decltype(auto) v_ref = to_ref(std::forward<Vec>(v));
+  const auto theta = (v_ref.array() - v_ref.maxCoeff()).exp();
+  return (theta / theta.sum()).matrix();
+}
+
+/**
+ * Return the softmax of each vector in an array.
+ *
+ * @tparam T `std::vector` whose scalar type is arithmetic
+ * @param[in] x Array of vectors to transform.
+ * @return Array of unit simplex results.
+ */
+template <typename T, require_std_vector_st<std::is_arithmetic, T>* = nullptr>
+inline auto softmax(T&& x) {
+  return apply_vector_unary<T>::apply(std::forward<T>(x), [](auto&& v) {
+    return softmax(std::forward<decltype(v)>(v));
+  });
 }
 
 }  // namespace math