refactor to support array-of-vectors instead of matrices + some other cleanup

Author: jachymb

stan/math/fwd/fun/log_softmax.hpp

@@ -5,72 +5,67 @@
 #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 rows of the specified matrix.
- * Each row is transformed independently; the result has the same shape
- * as the input.
+ * Return the log softmax of the specified row vector of `fvar` values.
+ * Delegates to the column-vector overload via transposition.
  *
- * @tparam Mat type of input matrix (Eigen matrix with fvar scalar)
- * @param[in] m Matrix to transform row-wise.
- * @return Log-softmax applied row-wise.
- * @throw std::domain_error If the input matrix is size 0.
+ * @tparam RowVec Eigen row vector with `fvar` scalar
+ * @param x row vector to transform
+ * @return log softmax of the row vector
  */
-template <typename Mat, require_eigen_t<Mat>* = nullptr,
-          require_not_eigen_vector_t<Mat>* = nullptr,
-          require_t<is_fvar<value_type_t<Mat>>>* = nullptr>
-inline auto log_softmax(const Mat& m) {
-  check_nonzero_size("log_softmax", "m", m);
-  const auto& m_ref = to_ref(m);
-  const auto val = m_ref.val().eval();
-  const auto shifted
-      = (val.array().colwise() - val.rowwise().maxCoeff().array()).eval();
-  const auto exp_s = shifted.exp().eval();
-  const auto row_sums = exp_s.rowwise().sum().eval();
-  const auto lsm_val = (shifted.colwise() - row_sums.log()).matrix().eval();
-  // softmax values needed for the tangent: d_in - softmax(x) ⊙ dot(softmax(x),
-  // d_in)
-  const auto s = (exp_s.colwise() / row_sums).eval();
-  const auto d_in = m_ref.d().eval();
-  const auto dots = (s.array() * d_in.array()).rowwise().sum().eval();
-  plain_type_t<Mat> result(m_ref.rows(), m_ref.cols());
-  result.val() = lsm_val;
-  result.d() = (d_in.array().colwise() - dots.array()).matrix();
-  return result;
-}
-
 template <typename RowVec, require_eigen_row_vector_t<RowVec>* = nullptr,
           require_t<is_fvar<value_type_t<RowVec>>>* = nullptr>
 inline auto log_softmax(const RowVec& x) {
   return log_softmax(x.transpose()).transpose().eval();
 }
 
-template <typename T, require_vector_st<is_fvar, T>* = nullptr,
-          require_not_t<is_eigen_row_vector<std::decay_t<T>>>* = nullptr>
+/**
+ * Return the log 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 log softmax results
+ */
+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 = std::decay_t<decltype(alpha)>;
-    using T_fvar = value_type_t<T_alpha>;
-    using T_inner = typename T_fvar::Scalar;
-
-    auto&& alpha_ref = to_ref(std::forward<decltype(alpha)>(alpha));
-    const Eigen::Matrix<T_inner, -1, 1> val = alpha_ref.val();
-    const Eigen::Matrix<T_inner, -1, 1> s = softmax(val);
-    const auto d_in = alpha_ref.d().eval();
-    const T_inner dot_sd = s.dot(d_in);
-
-    Eigen::Matrix<T_fvar, -1, 1> result(alpha_ref.size());
-    result.val() = log_softmax(val);
-    result.d() = (d_in.array() - dot_sd).matrix();
-    return result;
+  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 column vector of `fvar` values.
+ *
+ * @tparam ColVec Eigen column vector with `fvar` scalar
+ * @param x column vector to transform
+ * @return log softmax of the column vector
+ * @throw std::domain_error if the input size is 0
+ */
+template <typename ColVec,
+          require_eigen_col_vector_vt<is_fvar, ColVec>* = nullptr>
+inline auto log_softmax(const ColVec& x) {
+  using Eigen::Dynamic;
+  using Eigen::Matrix;
+  using T = typename value_type_t<ColVec>::Scalar;
+  check_nonzero_size("log_softmax", "x", x);
+  const auto& x_ref = to_ref(x);
+  const Matrix<T, Dynamic, 1> s = softmax(value_of(x_ref));
+  const auto d_in = x_ref.d().eval();
+  const T dot_sd = s.dot(d_in);
+  Matrix<fvar<T>, Dynamic, 1> 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,45 +6,60 @@
 #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 Mat, require_eigen_t<Mat>* = nullptr,
-          require_not_eigen_vector_t<Mat>* = nullptr,
-          require_t<is_fvar<value_type_t<Mat>>>* = nullptr>
-inline auto softmax(const Mat& m) {
-  const auto& m_ref = to_ref(m);
-  const auto s = softmax(m_ref.val());
-  const auto d_in = m_ref.d().eval();
-  // d/dx softmax(x) applied to tangent: s ⊙ (d_in - s · d_in)  (per row)
-  const auto dots = (s.array() * d_in.array()).rowwise().sum().eval();
-  plain_type_t<Mat> result(m_ref.rows(), m_ref.cols());
-  result.val() = s;
-  result.d() = (s.array() * (d_in.array().colwise() - dots.array())).matrix();
-  return result;
-}
-
+/**
+ * Return the softmax of the specified row vector of `fvar` values.
+ * Delegates to the column-vector overload via transposition.
+ *
+ * @tparam RowVec Eigen row vector with `fvar` scalar
+ * @param x row vector to transform
+ * @return softmax of the row vector
+ */
 template <typename RowVec, require_eigen_row_vector_t<RowVec>* = nullptr,
           require_t<is_fvar<value_type_t<RowVec>>>* = nullptr>
-inline auto softmax(const RowVec& alpha) {
-  return softmax(alpha.transpose()).transpose().eval();
+inline auto softmax(const RowVec& x) {
+  return softmax(x.transpose()).transpose().eval();
+}
+
+/**
+ * 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));
+  });
 }
 
+/**
+ * Return the softmax of the specified column vector of `fvar` values.
+ *
+ * @tparam ColVec Eigen column vector with `fvar` scalar
+ * @param x column vector to transform
+ * @return softmax of the column vector
+ */
 template <typename ColVec,
           require_eigen_col_vector_vt<is_fvar, ColVec>* = nullptr>
-inline auto softmax(const ColVec& alpha) {
+inline auto softmax(const ColVec& x) {
   using Eigen::Dynamic;
   using Eigen::Matrix;
   using T = typename value_type_t<ColVec>::Scalar;
-  if (alpha.size() == 0) {
+  if (x.size() == 0) {
     return Matrix<fvar<T>, Dynamic, 1>();
   }
-  const auto& alpha_ref = to_ref(alpha);
-  const Matrix<T, Dynamic, 1> s = softmax(value_of(alpha_ref));
-  const auto d_in = alpha_ref.d().eval();
+  const auto& x_ref = to_ref(x);
+  const Matrix<T, Dynamic, 1> s = softmax(value_of(x_ref));
+  const auto d_in = x_ref.d().eval();
   const T dot_sd = s.dot(d_in);
-  Matrix<fvar<T>, Dynamic, 1> result(alpha.size());
+  Matrix<fvar<T>, Dynamic, 1> result(x_ref.size());
   result.val() = s;
   result.d() = (s.array() * (d_in.array() - dot_sd)).matrix();
   return result;

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,29 +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_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(
@@ -55,26 +57,6 @@ inline auto log_softmax(Container&& x) {
       to_ref(std::forward<Container>(x)));
 }
 
-/**
- * Return the log softmax of the rows of the specified matrix.
- * Each row is transformed independently; the result has the same shape
- * as the input.
- *
- * @tparam Mat type of input matrix
- * @param[in] m Matrix to transform row-wise.
- * @return Log-softmax applied row-wise.
- */
-template <typename Mat, require_eigen_vt<std::is_arithmetic, Mat>* = nullptr,
-          require_not_eigen_vector_t<Mat>* = nullptr>
-inline plain_type_t<Mat> log_softmax(const Mat& m) {
-  check_nonzero_size("log_softmax", "m", m);
-  const auto& m_ref = to_ref(m);
-  const auto shifted
-      = (m_ref.array().colwise() - m_ref.rowwise().maxCoeff().array()).eval();
-  const auto exp_s = shifted.exp().eval();
-  return (shifted.colwise() - exp_s.rowwise().sum().log()).matrix();
-}
-
 }  // namespace math
 }  // namespace stan
 #endif

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,7 +39,7 @@ namespace math {
  * \end{array}
  * \f$
  *
- * @tparam Vec 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.
  */
@@ -54,22 +55,17 @@ inline plain_type_t<Vec> softmax(const Vec& v) {
 }
 
 /**
- * Return the softmax of the rows of the specified matrix.
- * Each row is transformed independently; the result is a row-stochastic
- * matrix whose rows each sum to one.
+ * Return the softmax of each vector in an array.
  *
- * @tparam Mat type of input matrix
- * @param[in] m Matrix to transform row-wise.
- * @return Row-stochastic matrix result of applying softmax to each row.
+ * @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 Mat, require_eigen_vt<std::is_arithmetic, Mat>* = nullptr,
-          require_not_eigen_vector_t<Mat>* = nullptr>
-inline plain_type_t<Mat> softmax(const Mat& m) {
-  const auto& m_ref = to_ref(m);
-  const auto shifted
-      = (m_ref.array().colwise() - m_ref.rowwise().maxCoeff().array()).eval();
-  const auto exp_s = shifted.exp().eval();
-  return (exp_s.colwise() / exp_s.rowwise().sum()).matrix();
+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