Vectorise binary log_sum_exp

Author: andrjohns

stan/math/prim/fun/log_sum_exp.hpp

@@ -5,6 +5,7 @@
 #include <stan/math/prim/fun/constants.hpp>
 #include <stan/math/prim/fun/Eigen.hpp>
 #include <stan/math/prim/fun/log1p_exp.hpp>
+#include <stan/math/prim/functor/apply_scalar_binary.hpp>
 #include <cmath>
 #include <vector>
 
@@ -47,7 +48,8 @@ namespace math {
  * @param a the first variable
  * @param b the second variable
  */
-template <typename T1, typename T2, require_all_not_st_var<T1, T2>* = nullptr>
+template <typename T1, typename T2, require_all_not_st_var<T1, T2>* = nullptr,
+          require_all_stan_scalar_t<T1, T2>* = nullptr>
 inline return_type_t<T1, T2> log_sum_exp(const T2& a, const T1& b) {
   if (a == NEGATIVE_INFTY) {
     return b;
@@ -91,6 +93,22 @@ inline auto log_sum_exp(const T& x) {
   });
 }
 
+/**
+ * Enables the vectorized application of the log_sum_exp function,
+ * when the first and/or second arguments are containers.
+ *
+ * @tparam T1 type of first input
+ * @tparam T2 type of second input
+ * @param a First input
+ * @param b Second input
+ * @return log_sum_exp function applied to the two inputs.
+ */
+template <typename T1, typename T2, require_any_container_t<T1, T2>* = nullptr>
+inline auto log_sum_exp(const T1& a, const T2& b) {
+  return apply_scalar_binary(
+      a, b, [](const auto& c, const auto& d) { return log_sum_exp(c, d); });
+}
+
 }  // namespace math
 }  // namespace stan
 

test/unit/math/mix/fun/log_sum_exp_test.cpp

@@ -96,3 +96,16 @@ TEST(MathMixMatFun, logSumExp) {
       std::vector<double>(x2c.data(), x2c.data() + x2c.size())};
   stan::test::expect_ad(tols, f, ststx);
 }
+
+TEST(mathMixScalFun, logSumExp_vec) {
+  auto f = [](const auto& x1, const auto& x2) {
+    using stan::math::log_sum_exp;
+    return log_sum_exp(x1, x2);
+  };
+
+  Eigen::VectorXd in1(2);
+  in1 << 3, 1;
+  Eigen::VectorXd in2(2);
+  in2 << 0.5, 3.4;
+  stan::test::expect_ad_vectorized_binary(f, in1, in2);
+}