Generalised unary vector function framework

stan/math/fwd/arr.hpp

@@ -7,7 +7,6 @@
 #include <stan/math/prim/arr.hpp>
 #include <stan/math/fwd/scal.hpp>
 
-#include <stan/math/fwd/arr/fun/log_sum_exp.hpp>
 #include <stan/math/fwd/arr/fun/sum.hpp>
 #include <stan/math/fwd/arr/fun/to_fvar.hpp>
 

stan/math/fwd/mat/fun/log_softmax.hpp

@@ -6,44 +6,47 @@
 #include <stan/math/prim/mat/fun/Eigen.hpp>
 #include <stan/math/prim/mat/fun/log_softmax.hpp>
 #include <stan/math/prim/mat/fun/softmax.hpp>
+#include <stan/math/prim/meta.hpp>
 
 namespace stan {
 namespace math {
 
-template <typename T>
-inline Eigen::Matrix<fvar<T>, Eigen::Dynamic, 1> log_softmax(
-    const Eigen::Matrix<fvar<T>, Eigen::Dynamic, 1>& alpha) {
-  using Eigen::Dynamic;
-  using Eigen::Matrix;
-
-  Matrix<T, Dynamic, 1> alpha_t(alpha.size());
-  for (int k = 0; k < alpha.size(); ++k) {
-    alpha_t(k) = alpha(k).val_;
-  }
-
-  Matrix<T, Dynamic, 1> softmax_alpha_t = softmax(alpha_t);
-  Matrix<T, Dynamic, 1> log_softmax_alpha_t = log_softmax(alpha_t);
-
-  Matrix<fvar<T>, Dynamic, 1> log_softmax_alpha(alpha.size());
-  for (int k = 0; k < alpha.size(); ++k) {
-    log_softmax_alpha(k).val_ = log_softmax_alpha_t(k);
-    log_softmax_alpha(k).d_ = 0;
-  }
-
-  for (int m = 0; m < alpha.size(); ++m) {
-    T negative_alpha_m_d_times_softmax_alpha_t_m
-        = -alpha(m).d_ * softmax_alpha_t(m);
-    for (int k = 0; k < alpha.size(); ++k) {
-      if (m == k) {
-        log_softmax_alpha(k).d_
-            += alpha(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 the log softmax of the specified vector or container of vectors.
+ *
+ * @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.
+ */
+template <typename T, require_t<is_fvar<scalar_type_t<T>>>...>
+inline auto log_softmax(T&& x) {
+  return apply_vector_unary<T>::apply(std::forward<T>(x), [&](auto& alpha) {
+    using T_fvar = value_type_t<decltype(alpha)>;
+    using T_fvar_inner = typename T_fvar::Scalar;
+
+    Eigen::Matrix<T_fvar_inner, -1, 1> alpha_t = alpha.val();
+    Eigen::Matrix<T_fvar_inner, -1, 1> softmax_alpha_t = softmax(alpha_t);
+
+    Eigen::Matrix<T_fvar, -1, 1> log_softmax_alpha(alpha.size());
+    log_softmax_alpha.val() = log_softmax(alpha_t);
+    log_softmax_alpha.d().setZero();
+
+    for (int m = 0; m < alpha.size(); ++m) {
+      T_fvar_inner negative_alpha_m_d_times_softmax_alpha_t_m
+          = -alpha(m).d_ * softmax_alpha_t(m);
+      for (int k = 0; k < alpha.size(); ++k) {
+        if (m == k) {
+          log_softmax_alpha(k).d_
+              += alpha(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 log_softmax_alpha;
+  });
 }
 
 }  // namespace math

stan/math/fwd/mat/fun/log_sum_exp.hpp

@@ -2,19 +2,39 @@
 #define STAN_MATH_FWD_MAT_FUN_LOG_SUM_EXP_HPP
 
 #include <stan/math/prim/mat/fun/Eigen.hpp>
+#include <stan/math/prim/mat/vectorize/apply_vector_unary.hpp>
 #include <stan/math/fwd/core.hpp>
 #include <stan/math/prim/mat/fun/log_sum_exp.hpp>
 
 namespace stan {
 namespace math {
 
-template <typename T, int R, int C>
-fvar<T> log_sum_exp(const Eigen::Matrix<fvar<T>, R, C>& v) {
-  Eigen::Matrix<T, R, C> vals = v.val();
-  Eigen::Matrix<T, R, C> exp_vals = vals.array().exp();
+/**
+ * Return the log of the sum of the exponentiated values of the specified
+ * matrix of values.  The matrix may be a full matrix, a vector,
+ * a row vector, or a container of these.
+ *
+ * The function is defined as follows to prevent overflow in exponential
+ * calculations.
+ *
+ * \f$\log \sum_{n=1}^N \exp(x_n) = \max(x) + \log \sum_{n=1}^N \exp(x_n -
+ * \max(x))\f$.
+ *
+ * @tparam T Type of input vector or matrix.
+ * @param[in] x Matrix of specified values.
+ * @return The log of the sum of the exponentiated vector values.
+ */
+template <typename T, require_t<is_fvar<scalar_type_t<T>>>...>
+inline auto log_sum_exp(T&& x) {
+  return apply_vector_unary<T>::reduce(std::forward<T>(x), [&](auto& v) {
+    using T_fvar_inner = typename value_type_t<decltype(v)>::Scalar;
+    using mat_type = Eigen::Matrix<T_fvar_inner, -1, -1>;
+    mat_type vals = v.val();
+    mat_type exp_vals = vals.array().exp();
 
-  return fvar<T>(log_sum_exp(vals),
-                 v.d().cwiseProduct(exp_vals).sum() / exp_vals.sum());
+    return fvar<T_fvar_inner>(
+        log_sum_exp(vals), v.d().cwiseProduct(exp_vals).sum() / exp_vals.sum());
+  });
 }
 
 }  // namespace math

stan/math/prim/arr.hpp

@@ -10,7 +10,6 @@
 #include <stan/math/prim/arr/fun/dot_self.hpp>
 #include <stan/math/prim/arr/fun/fill.hpp>
 #include <stan/math/prim/arr/fun/inverse_softmax.hpp>
-#include <stan/math/prim/arr/fun/log_sum_exp.hpp>
 #include <stan/math/prim/arr/fun/promote_elements.hpp>
 #include <stan/math/prim/arr/fun/promote_scalar.hpp>
 #include <stan/math/prim/arr/fun/promote_scalar_type.hpp>

stan/math/prim/mat.hpp

@@ -332,6 +332,7 @@
 #include <stan/math/prim/mat/prob/wishart_rng.hpp>
 
 #include <stan/math/prim/mat/vectorize/apply_scalar_unary.hpp>
+#include <stan/math/prim/mat/vectorize/apply_vector_unary.hpp>
 
 #include <stan/math/prim/arr.hpp>
 

stan/math/prim/mat/fun/log_softmax.hpp

@@ -4,6 +4,7 @@
 #include <stan/math/prim/err.hpp>
 #include <stan/math/prim/mat/fun/Eigen.hpp>
 #include <stan/math/prim/mat/fun/log_sum_exp.hpp>
+#include <stan/math/prim/mat/vectorize/apply_vector_unary.hpp>
 
 namespace stan {
 namespace math {
@@ -32,18 +33,17 @@ namespace math {
  * \right.
  * \f$
  *
- * @tparam T type of elements in the vector
- * @param[in] v Vector to transform.
- * @return Unit simplex result of the softmax transform of the vector.
+ * @tparam T Type of input vector to transform.
+ * @param[in] x Vector to transform.
+ * @return log unit simplex result of the softmax transform of the vector.
  */
-template <typename T>
-inline Eigen::Matrix<T, Eigen::Dynamic, 1> log_softmax(
-    const Eigen::Matrix<T, Eigen::Dynamic, 1>& v) {
-  check_nonzero_size("log_softmax", "v", v);
-  return v.array() - log_sum_exp(v);
+template <typename T, require_t<std::is_arithmetic<scalar_type_t<T>>>...>
+inline auto log_softmax(T&& x) {
+  return apply_vector_unary<T>::apply(std::forward<T>(x), [&](auto& v) {
+    check_nonzero_size("log_softmax", "v", v);
+    return (v.array() - log_sum_exp(v)).matrix();
+  });
 }
-
 }  // namespace math
 }  // namespace stan
-
 #endif

stan/math/prim/mat/fun/log_sum_exp.hpp

@@ -2,6 +2,7 @@
 #define STAN_MATH_PRIM_MAT_FUN_LOG_SUM_EXP_HPP
 
 #include <stan/math/prim/mat/fun/Eigen.hpp>
+#include <stan/math/prim/mat/vectorize/apply_vector_unary.hpp>
 #include <stan/math/prim/scal/fun/constants.hpp>
 #include <cmath>
 #include <vector>
@@ -12,31 +13,30 @@ namespace math {
 /**
  * Return the log of the sum of the exponentiated values of the specified
  * matrix of values.  The matrix may be a full matrix, a vector,
- * or a row vector.
+ * a row vector, or a container of these.
  *
  * The function is defined as follows to prevent overflow in exponential
  * calculations.
  *
  * \f$\log \sum_{n=1}^N \exp(x_n) = \max(x) + \log \sum_{n=1}^N \exp(x_n -
  * \max(x))\f$.
  *
- * @tparam R number of rows, can be Eigen::Dynamic
- * @tparam C number of columns, can be Eigen::Dynamic
- *
- * @param[in] x Matrix of specified values
+ * @tparam T Type of input vector or matrix.
+ * @param[in] x Matrix of specified values.
  * @return The log of the sum of the exponentiated vector values.
  */
-template <int R, int C>
-double log_sum_exp(const Eigen::Matrix<double, R, C>& x) {
-  if (x.size() == 0) {
-    return NEGATIVE_INFTY;
-  }
-
-  const double max = x.maxCoeff();
-  if (!std::isfinite(max)) {
-    return max;
-  }
-  return max + std::log((x.array() - max).exp().sum());
+template <typename T, require_t<std::is_arithmetic<scalar_type_t<T>>>...>
+inline auto log_sum_exp(T&& x) {
+  return apply_vector_unary<T>::reduce(std::forward<T>(x), [&](auto& v) {
+    if (v.size() == 0) {
+      return NEGATIVE_INFTY;
+    }
+    const double max = v.maxCoeff();
+    if (!std::isfinite(max)) {
+      return max;
+    }
+    return max + std::log((v.array() - max).exp().sum());
+  });
 }
 
 }  // namespace math

stan/math/prim/mat/prob/categorical_logit_lpmf.hpp

@@ -3,7 +3,6 @@
 
 #include <stan/math/prim/meta.hpp>
 #include <stan/math/prim/err.hpp>
-#include <stan/math/prim/arr/fun/log_sum_exp.hpp>
 #include <stan/math/prim/mat/fun/log_softmax.hpp>
 #include <stan/math/prim/mat/fun/log_sum_exp.hpp>
 #include <stan/math/prim/mat/fun/sum.hpp>