From 3cc61f552918458a14e342aad0e0d507e7d6fdaa Mon Sep 17 00:00:00 2001
From: Jouni Helske <jvhels@utu.fi>
Date: Mon, 24 Aug 2026 12:48:31 +0300
Subject: [PATCH 1/7] additional tests for log_softmax and log_sum_exp

---
 test/unit/math/fwd/fun/log_softmax_test.cpp | 112 ++++++++++++++++++++
 test/unit/math/fwd/fun/log_sum_exp_test.cpp |  42 ++++++++
 test/unit/math/rev/fun/log_softmax_test.cpp |  61 +++++++++++
 test/unit/math/rev/fun/log_sum_exp_test.cpp |  51 +++++++++
 4 files changed, 266 insertions(+)
 create mode 100644 test/unit/math/fwd/fun/log_softmax_test.cpp
 create mode 100644 test/unit/math/fwd/fun/log_sum_exp_test.cpp

diff --git a/test/unit/math/fwd/fun/log_softmax_test.cpp b/test/unit/math/fwd/fun/log_softmax_test.cpp
new file mode 100644
index 00000000000..94ca0ea6622
--- /dev/null
+++ b/test/unit/math/fwd/fun/log_softmax_test.cpp
@@ -0,0 +1,112 @@
+#include <stan/math/fwd.hpp>
+#include <gtest/gtest.h>
+
+#include <cmath>
+
+TEST(AgradFwd, log_softmax_large_dynamic_range) {
+  using stan::math::fvar;
+  using stan::math::log_softmax;
+  
+  Eigen::Matrix<fvar<double>, Eigen::Dynamic, 1> x(3);
+  
+  x(0) = fvar<double>(0.0, 1.0);
+  x(1) = fvar<double>(-100.0, 2.0);
+  x(2) = fvar<double>(-1000.0, 3.0);
+  
+  auto y = log_softmax(x);
+  
+  /*
+   * log_softmax(x) = x - log_sum_exp(x).
+   *
+   * exp(-1000) underflows to zero in double precision. Therefore:
+   *
+   *   log(softmax(x)[2])
+   *
+   * gives -inf while
+   *
+   *   x[2] - log_sum_exp(x)
+   *
+   * correctly retains the log-probability.
+   */
+  
+  EXPECT_TRUE(std::isfinite(y.val()(0)));
+  EXPECT_TRUE(std::isfinite(y.val()(1)));
+  EXPECT_TRUE(std::isfinite(y.val()(2)));
+  
+  EXPECT_DOUBLE_EQ(0.0, y.val()(0));
+  EXPECT_DOUBLE_EQ(-100.0, y.val()(1));
+  EXPECT_DOUBLE_EQ(-1000.0, y.val()(2));
+  
+  const double p1 = std::exp(-100.0);
+  const double dot_sd = 1.0 + 2.0 * p1;
+  
+  EXPECT_NEAR(1.0 - dot_sd, y(0).d(), 1e-12);
+  EXPECT_NEAR(2.0 - dot_sd, y(1).d(), 1e-12);
+  EXPECT_NEAR(3.0 - dot_sd, y(2).d(), 1e-12);
+}
+
+TEST(AgradFwd, log_softmax_many_small_probabilities) {
+  using stan::math::fvar;
+  using stan::math::log_softmax;
+  
+  constexpr int n = 501;
+  
+  Eigen::Matrix<fvar<double>, Eigen::Dynamic, 1> x(n);
+  
+  x(0) = fvar<double>(0.0, 1.0);
+  
+  for (int i = 1; i < n; ++i) {
+    x(i) = fvar<double>(-36.0, 2.0);
+  }
+  
+  const auto y = log_softmax(x);
+  
+  const double lse = std::log1p(500.0 * std::exp(-36.0));
+  EXPECT_NEAR(-lse, y(0).val(), 1e-14);
+  EXPECT_NEAR(-36.0 - lse, y(1).val(), 1e-14);
+  EXPECT_NEAR(-36.0 - lse, y(500).val(), 1e-14);
+  
+  for (int i = 0; i < n; ++i) {
+    EXPECT_TRUE(std::isfinite(y(i).val()));
+    EXPECT_TRUE(std::isfinite(y(i).d()));
+  }
+  
+  const double p_small = std::exp(-36.0) / (1.0 + 500.0 * std::exp(-36.0));
+  const double p_large = 1.0 / (1.0 + 500.0 * std::exp(-36.0));
+  
+  const double tangent = p_large + 500.0 * 2.0 * p_small;
+  
+  EXPECT_NEAR(1.0 - tangent, y(0).d(), 1e-12);
+  EXPECT_NEAR(2.0 - tangent, y(1).d(), 1e-12);
+}
+TEST(Agrad, log_softmax_positive_infinity_consistent_with_prim) {
+  using stan::math::fvar;
+  using stan::math::log_softmax;
+  
+  constexpr double inf = std::numeric_limits<double>::infinity();
+  
+  Eigen::Matrix<fvar<double>, Eigen::Dynamic, 1> x(4);
+  x(0) = fvar<double>(1.0, 1.0);
+  x(1) = fvar<double>(2.0, 2.0);
+  x(2) = fvar<double>(3.0, 3.0);
+  x(3) = fvar<double>(inf, 4.0);
+  
+  auto y = log_softmax(x);
+  
+  // The value is defined as
+  //
+  //   x - log_sum_exp(x).
+  //
+  // Since log_sum_exp([1, 2, 3, +Inf]) = +Inf:
+  //
+  //   1 - Inf = -Inf
+  //   2 - Inf = -Inf
+  //   3 - Inf = -Inf
+  //   Inf - Inf = NaN
+  //
+  
+  EXPECT_EQ(-inf, y(0).val());
+  EXPECT_EQ(-inf, y(1).val());
+  EXPECT_EQ(-inf, y(2).val());
+  EXPECT_TRUE(std::isnan(y(3).val()));
+}
\ No newline at end of file
diff --git a/test/unit/math/fwd/fun/log_sum_exp_test.cpp b/test/unit/math/fwd/fun/log_sum_exp_test.cpp
new file mode 100644
index 00000000000..8251018f52c
--- /dev/null
+++ b/test/unit/math/fwd/fun/log_sum_exp_test.cpp
@@ -0,0 +1,42 @@
+#include <stan/math/fwd.hpp>
+#include <gtest/gtest.h>
+
+#include <cmath>
+
+TEST(AgradFwdLogSumExp, derivative_does_not_overflow) {
+  using stan::math::fvar;
+  using stan::math::log_sum_exp;
+  
+  Eigen::Matrix<fvar<double>, Eigen::Dynamic, 1> x(3);
+  x(0) = fvar<double>(1000.0, 1.0);
+  x(1) = fvar<double>(1001.0, 2.0);
+  x(2) = fvar<double>(999.0, 3.0);
+  
+  auto y = log_sum_exp(x);
+  // log_sum_exp([1000, 1001, 999])
+  //
+  // = 1001 + log(exp(-1) + 1 + exp(-2)).
+  const double expected_value
+      = 1001.0 + std::log(std::exp(-1.0) + 1.0 + std::exp(-2.0));
+  
+  EXPECT_NEAR(expected_value, y.val(), 1e-12);
+  EXPECT_TRUE(std::isfinite(y.val()));
+  
+  // The derivative is
+  //
+  // softmax(x) dot x.d().
+  //
+  // The old implementation formed exp(x) directly in the derivative,
+  // which gives inf / inf for these values.
+  
+  const double denom = std::exp(-1.0) + 1.0 + std::exp(-2.0);
+  
+  const double p0 = std::exp(-1.0) / denom;
+  const double p1 = 1.0 / denom;
+  const double p2 = std::exp(-2.0) / denom;
+  
+  const double expected_derivative = p0 + 2.0 * p1 + 3.0 * p2;
+  
+  EXPECT_TRUE(std::isfinite(y.d()));
+  EXPECT_NEAR(expected_derivative, y.d(), 1e-12);
+}
diff --git a/test/unit/math/rev/fun/log_softmax_test.cpp b/test/unit/math/rev/fun/log_softmax_test.cpp
index 0008e372b4c..b5d97698854 100644
--- a/test/unit/math/rev/fun/log_softmax_test.cpp
+++ b/test/unit/math/rev/fun/log_softmax_test.cpp
@@ -60,3 +60,64 @@ TEST_F(AgradRev, log_softmax_var_value_row_vector) {
   EXPECT_FLOAT_EQ(1.0 - 3.0 * std::exp(1.0) / denom, x.adj()(1));
   EXPECT_FLOAT_EQ(1.0 - 3.0 * std::exp(2.0) / denom, x.adj()(2));
 }
+
+TEST(AgradRev, log_softmax_negative_infinity_has_finite_adjoint) {
+  using stan::math::log_softmax;
+  using stan::math::sum;
+  using stan::math::var_value;
+  
+  const double neg_inf = -std::numeric_limits<double>::infinity();
+  
+  Eigen::VectorXd v(4);
+  v << neg_inf, 1.0, 2.0, 3.0;
+  
+  var_value<Eigen::VectorXd> x(v);
+  
+  const auto y = log_softmax(x);
+  sum(y).grad();
+  
+  const double e1 = std::exp(1.0);
+  const double e2 = std::exp(2.0);
+  const double e3 = std::exp(3.0);
+  const double denom = e1 + e2 + e3;
+  
+  const double p0 = 0.0;
+  const double p1 = e1 / denom;
+  const double p2 = e2 / denom;
+  const double p3 = e3 / denom;
+  
+  EXPECT_NEAR(1.0, x.adj()(0), 1e-12);
+  EXPECT_NEAR(1.0 - 4.0 * p1, x.adj()(1), 1e-12);
+  EXPECT_NEAR(1.0 - 4.0 * p2, x.adj()(2), 1e-12);
+  EXPECT_NEAR(1.0 - 4.0 * p3, x.adj()(3), 1e-12);
+}
+
+TEST(AgradRev, log_softmax_adjoint_uses_stable_softmax) {
+  using stan::math::log_softmax;
+  using stan::math::sum;
+  using stan::math::var_value;
+  
+  Eigen::VectorXd v(4);
+  
+  v << 629.7901581243797,
+       31.52411463,
+       608.19720553,
+       120.94829574;
+  
+  var_value<Eigen::VectorXd> x(v);
+  auto y = log_softmax(x);
+  
+  sum(y).grad();
+  
+  // Independent max-shifted softmax calculation
+  Eigen::VectorXd p = (v.array() -  v.maxCoeff()).exp();
+  p /= p.sum();
+  
+  // d/dx_m sum_k log_softmax(x)_k
+  // = 1 - n * softmax(x)_m.
+  for (Eigen::Index i = 0; i < v.size(); ++i) {
+    const double expected = 1.0 - v.size() * p(i);
+    
+    EXPECT_NEAR(expected, x.adj()(i), 1e-12);
+  }
+}
diff --git a/test/unit/math/rev/fun/log_sum_exp_test.cpp b/test/unit/math/rev/fun/log_sum_exp_test.cpp
index f020a0d1dd2..7e7fdbadc26 100644
--- a/test/unit/math/rev/fun/log_sum_exp_test.cpp
+++ b/test/unit/math/rev/fun/log_sum_exp_test.cpp
@@ -61,3 +61,54 @@ TEST_F(AgradRev, log_sum_exp_tests_large_values) {
   output5.grad();
   EXPECT_FLOAT_EQ(a8.adj(), 0.0);
 }
+
+TEST(AgradRev, log_sum_exp_negative_infinity_has_zero_adjoint) {
+  using stan::math::log_sum_exp;
+  using stan::math::var_value;
+  
+  const double neg_inf = -std::numeric_limits<double>::infinity();
+  
+  Eigen::VectorXd v(4);
+  v << neg_inf, 1.0, 2.0, 3.0;
+  
+  var_value<Eigen::VectorXd> x(v);
+  
+  const auto y = log_sum_exp(x);
+  y.grad();
+  
+  const double e1 = std::exp(1.0);
+  const double e2 = std::exp(2.0);
+  const double e3 = std::exp(3.0);
+  const double denom = e1 + e2 + e3;
+  
+  EXPECT_DOUBLE_EQ(0.0, x.adj()(0));
+  EXPECT_NEAR(e1 / denom, x.adj()(1), 1e-12);
+  EXPECT_NEAR(e2 / denom, x.adj()(2), 1e-12);
+  EXPECT_NEAR(e3 / denom, x.adj()(3), 1e-12);
+}
+
+TEST(AgradRev, log_sum_exp_adjoint_uses_stable_softmax) {
+  using stan::math::log_sum_exp;
+  using stan::math::var_value;
+  
+  Eigen::VectorXd v(4);
+  
+  v << 629.7901581243797,
+       31.52411463,
+       608.19720553,
+       120.94829574;
+  
+  var_value<Eigen::VectorXd> x(v);
+  
+  auto y = log_sum_exp(x);
+  y.grad();
+  
+  // Independent max-shifted softmax calculation
+  Eigen::VectorXd p = (v.array() - v.maxCoeff()).exp();
+  p /= p.sum();
+  
+  // The gradient of log_sum_exp is softmax.
+  for (Eigen::Index i = 0; i < v.size(); ++i) {
+    EXPECT_NEAR(p(i), x.adj()(i), 1e-12);
+  }
+}

From 873012f60bcec473fefe2521605fb585cc6e69e1 Mon Sep 17 00:00:00 2001
From: Jouni Helske <jvhels@utu.fi>
Date: Mon, 24 Aug 2026 13:06:27 +0300
Subject: [PATCH 2/7] fix fwd tests (remove const)

---
 test/unit/math/rev/fun/log_softmax_test.cpp | 2 +-
 test/unit/math/rev/fun/log_sum_exp_test.cpp | 2 +-
 2 files changed, 2 insertions(+), 2 deletions(-)

diff --git a/test/unit/math/rev/fun/log_softmax_test.cpp b/test/unit/math/rev/fun/log_softmax_test.cpp
index b5d97698854..93d813d77fe 100644
--- a/test/unit/math/rev/fun/log_softmax_test.cpp
+++ b/test/unit/math/rev/fun/log_softmax_test.cpp
@@ -73,7 +73,7 @@ TEST(AgradRev, log_softmax_negative_infinity_has_finite_adjoint) {
   
   var_value<Eigen::VectorXd> x(v);
   
-  const auto y = log_softmax(x);
+  auto y = log_softmax(x);
   sum(y).grad();
   
   const double e1 = std::exp(1.0);
diff --git a/test/unit/math/rev/fun/log_sum_exp_test.cpp b/test/unit/math/rev/fun/log_sum_exp_test.cpp
index 7e7fdbadc26..1baa4354878 100644
--- a/test/unit/math/rev/fun/log_sum_exp_test.cpp
+++ b/test/unit/math/rev/fun/log_sum_exp_test.cpp
@@ -73,7 +73,7 @@ TEST(AgradRev, log_sum_exp_negative_infinity_has_zero_adjoint) {
   
   var_value<Eigen::VectorXd> x(v);
   
-  const auto y = log_sum_exp(x);
+  auto y = log_sum_exp(x);
   y.grad();
   
   const double e1 = std::exp(1.0);

From 5af9cede32d5dcbb79e2caf44bb9ac015f652769 Mon Sep 17 00:00:00 2001
From: Jouni Helske <jvhels@utu.fi>
Date: Mon, 24 Aug 2026 13:49:59 +0300
Subject: [PATCH 3/7] TEST to TEST_F

---
 test/unit/math/fwd/fun/log_softmax_test.cpp | 2 +-
 test/unit/math/fwd/fun/log_sum_exp_test.cpp | 2 +-
 test/unit/math/rev/fun/log_softmax_test.cpp | 4 ++--
 test/unit/math/rev/fun/log_sum_exp_test.cpp | 4 ++--
 4 files changed, 6 insertions(+), 6 deletions(-)

diff --git a/test/unit/math/fwd/fun/log_softmax_test.cpp b/test/unit/math/fwd/fun/log_softmax_test.cpp
index 94ca0ea6622..9b6b014e191 100644
--- a/test/unit/math/fwd/fun/log_softmax_test.cpp
+++ b/test/unit/math/fwd/fun/log_softmax_test.cpp
@@ -79,7 +79,7 @@ TEST(AgradFwd, log_softmax_many_small_probabilities) {
   EXPECT_NEAR(1.0 - tangent, y(0).d(), 1e-12);
   EXPECT_NEAR(2.0 - tangent, y(1).d(), 1e-12);
 }
-TEST(Agrad, log_softmax_positive_infinity_consistent_with_prim) {
+TEST(AgradFwd, log_softmax_positive_infinity_consistent_with_prim) {
   using stan::math::fvar;
   using stan::math::log_softmax;
   
diff --git a/test/unit/math/fwd/fun/log_sum_exp_test.cpp b/test/unit/math/fwd/fun/log_sum_exp_test.cpp
index 8251018f52c..aea936284aa 100644
--- a/test/unit/math/fwd/fun/log_sum_exp_test.cpp
+++ b/test/unit/math/fwd/fun/log_sum_exp_test.cpp
@@ -3,7 +3,7 @@
 
 #include <cmath>
 
-TEST(AgradFwdLogSumExp, derivative_does_not_overflow) {
+TEST(AgradFwd, log_sum_exp_derivative_does_not_overflow) {
   using stan::math::fvar;
   using stan::math::log_sum_exp;
   
diff --git a/test/unit/math/rev/fun/log_softmax_test.cpp b/test/unit/math/rev/fun/log_softmax_test.cpp
index 93d813d77fe..bad9fcf4f75 100644
--- a/test/unit/math/rev/fun/log_softmax_test.cpp
+++ b/test/unit/math/rev/fun/log_softmax_test.cpp
@@ -61,7 +61,7 @@ TEST_F(AgradRev, log_softmax_var_value_row_vector) {
   EXPECT_FLOAT_EQ(1.0 - 3.0 * std::exp(2.0) / denom, x.adj()(2));
 }
 
-TEST(AgradRev, log_softmax_negative_infinity_has_finite_adjoint) {
+TEST_F(AgradRev, log_softmax_negative_infinity_has_finite_adjoint) {
   using stan::math::log_softmax;
   using stan::math::sum;
   using stan::math::var_value;
@@ -92,7 +92,7 @@ TEST(AgradRev, log_softmax_negative_infinity_has_finite_adjoint) {
   EXPECT_NEAR(1.0 - 4.0 * p3, x.adj()(3), 1e-12);
 }
 
-TEST(AgradRev, log_softmax_adjoint_uses_stable_softmax) {
+TEST_F(AgradRev, log_softmax_adjoint_uses_stable_softmax) {
   using stan::math::log_softmax;
   using stan::math::sum;
   using stan::math::var_value;
diff --git a/test/unit/math/rev/fun/log_sum_exp_test.cpp b/test/unit/math/rev/fun/log_sum_exp_test.cpp
index 1baa4354878..19dc1ae386e 100644
--- a/test/unit/math/rev/fun/log_sum_exp_test.cpp
+++ b/test/unit/math/rev/fun/log_sum_exp_test.cpp
@@ -62,7 +62,7 @@ TEST_F(AgradRev, log_sum_exp_tests_large_values) {
   EXPECT_FLOAT_EQ(a8.adj(), 0.0);
 }
 
-TEST(AgradRev, log_sum_exp_negative_infinity_has_zero_adjoint) {
+TEST_F(AgradRev, log_sum_exp_negative_infinity_has_zero_adjoint) {
   using stan::math::log_sum_exp;
   using stan::math::var_value;
   
@@ -87,7 +87,7 @@ TEST(AgradRev, log_sum_exp_negative_infinity_has_zero_adjoint) {
   EXPECT_NEAR(e3 / denom, x.adj()(3), 1e-12);
 }
 
-TEST(AgradRev, log_sum_exp_adjoint_uses_stable_softmax) {
+TEST_F(AgradRev, log_sum_exp_adjoint_uses_stable_softmax) {
   using stan::math::log_sum_exp;
   using stan::math::var_value;
   

From b99d75ad72ba919a29807d2b2230df3c2e3fd38e Mon Sep 17 00:00:00 2001
From: Stan Jenkins <mc.stanislaw@gmail.com>
Date: Mon, 24 Aug 2026 11:21:15 -0400
Subject: [PATCH 4/7] [Jenkins] auto-formatting by clang-format version
 10.0.0-4ubuntu1

---
 test/unit/math/fwd/fun/log_softmax_test.cpp | 46 ++++++++++-----------
 test/unit/math/fwd/fun/log_sum_exp_test.cpp | 16 +++----
 test/unit/math/rev/fun/log_softmax_test.cpp | 35 +++++++---------
 test/unit/math/rev/fun/log_sum_exp_test.cpp | 29 ++++++-------
 4 files changed, 60 insertions(+), 66 deletions(-)

diff --git a/test/unit/math/fwd/fun/log_softmax_test.cpp b/test/unit/math/fwd/fun/log_softmax_test.cpp
index 9b6b014e191..9055b896af2 100644
--- a/test/unit/math/fwd/fun/log_softmax_test.cpp
+++ b/test/unit/math/fwd/fun/log_softmax_test.cpp
@@ -6,15 +6,15 @@
 TEST(AgradFwd, log_softmax_large_dynamic_range) {
   using stan::math::fvar;
   using stan::math::log_softmax;
-  
+
   Eigen::Matrix<fvar<double>, Eigen::Dynamic, 1> x(3);
-  
+
   x(0) = fvar<double>(0.0, 1.0);
   x(1) = fvar<double>(-100.0, 2.0);
   x(2) = fvar<double>(-1000.0, 3.0);
-  
+
   auto y = log_softmax(x);
-  
+
   /*
    * log_softmax(x) = x - log_sum_exp(x).
    *
@@ -28,18 +28,18 @@ TEST(AgradFwd, log_softmax_large_dynamic_range) {
    *
    * correctly retains the log-probability.
    */
-  
+
   EXPECT_TRUE(std::isfinite(y.val()(0)));
   EXPECT_TRUE(std::isfinite(y.val()(1)));
   EXPECT_TRUE(std::isfinite(y.val()(2)));
-  
+
   EXPECT_DOUBLE_EQ(0.0, y.val()(0));
   EXPECT_DOUBLE_EQ(-100.0, y.val()(1));
   EXPECT_DOUBLE_EQ(-1000.0, y.val()(2));
-  
+
   const double p1 = std::exp(-100.0);
   const double dot_sd = 1.0 + 2.0 * p1;
-  
+
   EXPECT_NEAR(1.0 - dot_sd, y(0).d(), 1e-12);
   EXPECT_NEAR(2.0 - dot_sd, y(1).d(), 1e-12);
   EXPECT_NEAR(3.0 - dot_sd, y(2).d(), 1e-12);
@@ -48,51 +48,51 @@ TEST(AgradFwd, log_softmax_large_dynamic_range) {
 TEST(AgradFwd, log_softmax_many_small_probabilities) {
   using stan::math::fvar;
   using stan::math::log_softmax;
-  
+
   constexpr int n = 501;
-  
+
   Eigen::Matrix<fvar<double>, Eigen::Dynamic, 1> x(n);
-  
+
   x(0) = fvar<double>(0.0, 1.0);
-  
+
   for (int i = 1; i < n; ++i) {
     x(i) = fvar<double>(-36.0, 2.0);
   }
-  
+
   const auto y = log_softmax(x);
-  
+
   const double lse = std::log1p(500.0 * std::exp(-36.0));
   EXPECT_NEAR(-lse, y(0).val(), 1e-14);
   EXPECT_NEAR(-36.0 - lse, y(1).val(), 1e-14);
   EXPECT_NEAR(-36.0 - lse, y(500).val(), 1e-14);
-  
+
   for (int i = 0; i < n; ++i) {
     EXPECT_TRUE(std::isfinite(y(i).val()));
     EXPECT_TRUE(std::isfinite(y(i).d()));
   }
-  
+
   const double p_small = std::exp(-36.0) / (1.0 + 500.0 * std::exp(-36.0));
   const double p_large = 1.0 / (1.0 + 500.0 * std::exp(-36.0));
-  
+
   const double tangent = p_large + 500.0 * 2.0 * p_small;
-  
+
   EXPECT_NEAR(1.0 - tangent, y(0).d(), 1e-12);
   EXPECT_NEAR(2.0 - tangent, y(1).d(), 1e-12);
 }
 TEST(AgradFwd, log_softmax_positive_infinity_consistent_with_prim) {
   using stan::math::fvar;
   using stan::math::log_softmax;
-  
+
   constexpr double inf = std::numeric_limits<double>::infinity();
-  
+
   Eigen::Matrix<fvar<double>, Eigen::Dynamic, 1> x(4);
   x(0) = fvar<double>(1.0, 1.0);
   x(1) = fvar<double>(2.0, 2.0);
   x(2) = fvar<double>(3.0, 3.0);
   x(3) = fvar<double>(inf, 4.0);
-  
+
   auto y = log_softmax(x);
-  
+
   // The value is defined as
   //
   //   x - log_sum_exp(x).
@@ -104,7 +104,7 @@ TEST(AgradFwd, log_softmax_positive_infinity_consistent_with_prim) {
   //   3 - Inf = -Inf
   //   Inf - Inf = NaN
   //
-  
+
   EXPECT_EQ(-inf, y(0).val());
   EXPECT_EQ(-inf, y(1).val());
   EXPECT_EQ(-inf, y(2).val());
diff --git a/test/unit/math/fwd/fun/log_sum_exp_test.cpp b/test/unit/math/fwd/fun/log_sum_exp_test.cpp
index aea936284aa..e6286602bb5 100644
--- a/test/unit/math/fwd/fun/log_sum_exp_test.cpp
+++ b/test/unit/math/fwd/fun/log_sum_exp_test.cpp
@@ -6,37 +6,37 @@
 TEST(AgradFwd, log_sum_exp_derivative_does_not_overflow) {
   using stan::math::fvar;
   using stan::math::log_sum_exp;
-  
+
   Eigen::Matrix<fvar<double>, Eigen::Dynamic, 1> x(3);
   x(0) = fvar<double>(1000.0, 1.0);
   x(1) = fvar<double>(1001.0, 2.0);
   x(2) = fvar<double>(999.0, 3.0);
-  
+
   auto y = log_sum_exp(x);
   // log_sum_exp([1000, 1001, 999])
   //
   // = 1001 + log(exp(-1) + 1 + exp(-2)).
   const double expected_value
       = 1001.0 + std::log(std::exp(-1.0) + 1.0 + std::exp(-2.0));
-  
+
   EXPECT_NEAR(expected_value, y.val(), 1e-12);
   EXPECT_TRUE(std::isfinite(y.val()));
-  
+
   // The derivative is
   //
   // softmax(x) dot x.d().
   //
   // The old implementation formed exp(x) directly in the derivative,
   // which gives inf / inf for these values.
-  
+
   const double denom = std::exp(-1.0) + 1.0 + std::exp(-2.0);
-  
+
   const double p0 = std::exp(-1.0) / denom;
   const double p1 = 1.0 / denom;
   const double p2 = std::exp(-2.0) / denom;
-  
+
   const double expected_derivative = p0 + 2.0 * p1 + 3.0 * p2;
-  
+
   EXPECT_TRUE(std::isfinite(y.d()));
   EXPECT_NEAR(expected_derivative, y.d(), 1e-12);
 }
diff --git a/test/unit/math/rev/fun/log_softmax_test.cpp b/test/unit/math/rev/fun/log_softmax_test.cpp
index bad9fcf4f75..e847af5e04a 100644
--- a/test/unit/math/rev/fun/log_softmax_test.cpp
+++ b/test/unit/math/rev/fun/log_softmax_test.cpp
@@ -65,27 +65,27 @@ TEST_F(AgradRev, log_softmax_negative_infinity_has_finite_adjoint) {
   using stan::math::log_softmax;
   using stan::math::sum;
   using stan::math::var_value;
-  
+
   const double neg_inf = -std::numeric_limits<double>::infinity();
-  
+
   Eigen::VectorXd v(4);
   v << neg_inf, 1.0, 2.0, 3.0;
-  
+
   var_value<Eigen::VectorXd> x(v);
-  
+
   auto y = log_softmax(x);
   sum(y).grad();
-  
+
   const double e1 = std::exp(1.0);
   const double e2 = std::exp(2.0);
   const double e3 = std::exp(3.0);
   const double denom = e1 + e2 + e3;
-  
+
   const double p0 = 0.0;
   const double p1 = e1 / denom;
   const double p2 = e2 / denom;
   const double p3 = e3 / denom;
-  
+
   EXPECT_NEAR(1.0, x.adj()(0), 1e-12);
   EXPECT_NEAR(1.0 - 4.0 * p1, x.adj()(1), 1e-12);
   EXPECT_NEAR(1.0 - 4.0 * p2, x.adj()(2), 1e-12);
@@ -96,28 +96,25 @@ TEST_F(AgradRev, log_softmax_adjoint_uses_stable_softmax) {
   using stan::math::log_softmax;
   using stan::math::sum;
   using stan::math::var_value;
-  
+
   Eigen::VectorXd v(4);
-  
-  v << 629.7901581243797,
-       31.52411463,
-       608.19720553,
-       120.94829574;
-  
+
+  v << 629.7901581243797, 31.52411463, 608.19720553, 120.94829574;
+
   var_value<Eigen::VectorXd> x(v);
   auto y = log_softmax(x);
-  
+
   sum(y).grad();
-  
+
   // Independent max-shifted softmax calculation
-  Eigen::VectorXd p = (v.array() -  v.maxCoeff()).exp();
+  Eigen::VectorXd p = (v.array() - v.maxCoeff()).exp();
   p /= p.sum();
-  
+
   // d/dx_m sum_k log_softmax(x)_k
   // = 1 - n * softmax(x)_m.
   for (Eigen::Index i = 0; i < v.size(); ++i) {
     const double expected = 1.0 - v.size() * p(i);
-    
+
     EXPECT_NEAR(expected, x.adj()(i), 1e-12);
   }
 }
diff --git a/test/unit/math/rev/fun/log_sum_exp_test.cpp b/test/unit/math/rev/fun/log_sum_exp_test.cpp
index 19dc1ae386e..c5f40e51e5c 100644
--- a/test/unit/math/rev/fun/log_sum_exp_test.cpp
+++ b/test/unit/math/rev/fun/log_sum_exp_test.cpp
@@ -65,22 +65,22 @@ TEST_F(AgradRev, log_sum_exp_tests_large_values) {
 TEST_F(AgradRev, log_sum_exp_negative_infinity_has_zero_adjoint) {
   using stan::math::log_sum_exp;
   using stan::math::var_value;
-  
+
   const double neg_inf = -std::numeric_limits<double>::infinity();
-  
+
   Eigen::VectorXd v(4);
   v << neg_inf, 1.0, 2.0, 3.0;
-  
+
   var_value<Eigen::VectorXd> x(v);
-  
+
   auto y = log_sum_exp(x);
   y.grad();
-  
+
   const double e1 = std::exp(1.0);
   const double e2 = std::exp(2.0);
   const double e3 = std::exp(3.0);
   const double denom = e1 + e2 + e3;
-  
+
   EXPECT_DOUBLE_EQ(0.0, x.adj()(0));
   EXPECT_NEAR(e1 / denom, x.adj()(1), 1e-12);
   EXPECT_NEAR(e2 / denom, x.adj()(2), 1e-12);
@@ -90,23 +90,20 @@ TEST_F(AgradRev, log_sum_exp_negative_infinity_has_zero_adjoint) {
 TEST_F(AgradRev, log_sum_exp_adjoint_uses_stable_softmax) {
   using stan::math::log_sum_exp;
   using stan::math::var_value;
-  
+
   Eigen::VectorXd v(4);
-  
-  v << 629.7901581243797,
-       31.52411463,
-       608.19720553,
-       120.94829574;
-  
+
+  v << 629.7901581243797, 31.52411463, 608.19720553, 120.94829574;
+
   var_value<Eigen::VectorXd> x(v);
-  
+
   auto y = log_sum_exp(x);
   y.grad();
-  
+
   // Independent max-shifted softmax calculation
   Eigen::VectorXd p = (v.array() - v.maxCoeff()).exp();
   p /= p.sum();
-  
+
   // The gradient of log_sum_exp is softmax.
   for (Eigen::Index i = 0; i < v.size(); ++i) {
     EXPECT_NEAR(p(i), x.adj()(i), 1e-12);

From 14117777d7092406c0b04019c5cf6ec2498c2af3 Mon Sep 17 00:00:00 2001
From: Jouni Helske <jvhels@utu.fi>
Date: Mon, 24 Aug 2026 20:35:45 +0300
Subject: [PATCH 5/7] improve accuracy and consistency of softmax et al

---
 stan/math/fwd/fun/log_softmax.hpp  | 41 +++++++++++++++++-------------
 stan/math/fwd/fun/log_sum_exp.hpp  |  9 ++++---
 stan/math/fwd/fun/softmax.hpp      | 25 +++++++++---------
 stan/math/prim/fun/log_softmax.hpp | 23 ++++++++---------
 stan/math/prim/fun/softmax.hpp     | 18 ++++++-------
 stan/math/rev/fun/log_softmax.hpp  |  9 ++++---
 stan/math/rev/fun/log_sum_exp.hpp  |  6 ++---
 stan/math/rev/fun/softmax.hpp      | 12 +++++----
 8 files changed, 77 insertions(+), 66 deletions(-)

diff --git a/stan/math/fwd/fun/log_softmax.hpp b/stan/math/fwd/fun/log_softmax.hpp
index a9685d7ca00..24c96922b64 100644
--- a/stan/math/fwd/fun/log_softmax.hpp
+++ b/stan/math/fwd/fun/log_softmax.hpp
@@ -14,10 +14,11 @@ namespace stan {
 namespace math {
 
 /**
- * Return the log softmax of each vector in a container of `fvar` values.
+ * Return the log softmax of each vector or matrix in a container of `fvar` 
+ * values.
  *
  * @tparam T `std::vector` whose scalar type is `fvar`
- * @param x container of vectors to transform
+ * @param x container of vectors or matrices to transform
  * @return container of log softmax results
  */
 template <typename T, require_std_vector_st<is_fvar, T>* = nullptr>
@@ -28,31 +29,35 @@ inline auto log_softmax(T&& x) {
 }
 
 /**
- * Return the log softmax of the specified vector of `fvar` values.
+ * Return the log softmax of the specified vector or matrix of `fvar` values.
  *
- * @tparam Vec Eigen vector with `fvar` scalar
- * @param x vector to transform
- * @return log softmax of the vector, or an empty result if the input is empty
+ * @tparam Mat Eigen vector or matrix with `fvar` scalar
+ * @param x vector or matrix to transform
+ * @return log softmax of the vector or matrix, or an empty result if the 
+ * input is empty
  */
-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;
-  decltype(auto) x_ref = to_ref(std::forward<Vec>(x));
+template <typename Mat, require_eigen_vt<is_fvar, Mat>* = nullptr>
+inline auto log_softmax(Mat&& x) {
+  using mat = std::decay_t<Mat>;
+  constexpr int Rows = mat::RowsAtCompileTime;
+  constexpr int Cols = mat::ColsAtCompileTime;
+  using T = typename value_type_t<mat>::Scalar;
+  decltype(auto) x_ref = to_ref(std::forward<Mat>(x));
   if (x_ref.size() == 0) {
     return Eigen::Matrix<fvar<T>, Rows, Cols>{};
   }
-  const auto s = softmax(value_of(x_ref));
+  const auto x_val = value_of(x_ref);
+  const auto lse = log_sum_exp(x_val);
+  const auto s = softmax(x_val);
   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();
+  const auto dot_sd = (s.array() * d_in.array()).sum();
+  
+  Eigen::Matrix<fvar<T>, Rows, Cols> result(x_ref.rows(), x_ref.cols());
+  result.val() = (x_val.array() - lse).matrix();
   result.d() = (d_in.array() - dot_sd).matrix();
   return result;
 }
 
 }  // namespace math
 }  // namespace stan
-#endif
+#endif
\ No newline at end of file
diff --git a/stan/math/fwd/fun/log_sum_exp.hpp b/stan/math/fwd/fun/log_sum_exp.hpp
index 926838d143c..399b13e877e 100644
--- a/stan/math/fwd/fun/log_sum_exp.hpp
+++ b/stan/math/fwd/fun/log_sum_exp.hpp
@@ -8,6 +8,7 @@
 #include <stan/math/prim/fun/constants.hpp>
 #include <stan/math/prim/fun/to_ref.hpp>
 #include <stan/math/prim/fun/log_sum_exp.hpp>
+#include <stan/math/prim/fun/softmax.hpp>
 #include <cmath>
 #include <vector>
 
@@ -56,11 +57,11 @@ inline auto log_sum_exp(T&& x) {
         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();
-
+        
+        const auto probs = softmax(vals);
         return fvar<T_fvar_inner>(
-            log_sum_exp(vals),
-            v.d().cwiseProduct(exp_vals).sum() / exp_vals.sum());
+          log_sum_exp(vals),
+          v.d().cwiseProduct(probs).sum());
       });
 }
 
diff --git a/stan/math/fwd/fun/softmax.hpp b/stan/math/fwd/fun/softmax.hpp
index 97d019b6c23..c77cb35bf33 100644
--- a/stan/math/fwd/fun/softmax.hpp
+++ b/stan/math/fwd/fun/softmax.hpp
@@ -29,24 +29,25 @@ inline auto softmax(T&& x) {
 /**
  * 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, or an empty result if the input is empty
+ * @tparam Mat Eigen vector or matrix with `fvar` scalar
+ * @param x vector or matrix to transform
+ * @return softmax of the vector, matrix, or an empty result if the input is 
+ *   empty
  */
-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;
-  decltype(auto) x_ref = to_ref(std::forward<Vec>(x));
+template <typename Mat, require_eigen_vt<is_fvar, Mat>* = nullptr>
+inline auto softmax(Mat&& x) {
+  using mat = std::decay_t<Mat>;
+  constexpr int Rows = mat::RowsAtCompileTime;
+  constexpr int Cols = mat::ColsAtCompileTime;
+  using T = typename value_type_t<mat>::Scalar;
+  decltype(auto) x_ref = to_ref(std::forward<Mat>(x));
   if (x_ref.size() == 0) {
     return Eigen::Matrix<fvar<T>, Rows, Cols>{};
   }
   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());
+  const auto dot_sd = (s.array() * d_in.array()).sum();
+  Eigen::Matrix<fvar<T>, Rows, Cols> result(x_ref.rows(), x_ref.cols());
   result.val() = s;
   result.d() = (s.array() * (d_in.array() - dot_sd)).matrix();
   return result;
diff --git a/stan/math/prim/fun/log_softmax.hpp b/stan/math/prim/fun/log_softmax.hpp
index d6c9017edb6..51665b47adb 100644
--- a/stan/math/prim/fun/log_softmax.hpp
+++ b/stan/math/prim/fun/log_softmax.hpp
@@ -12,9 +12,10 @@ namespace stan {
 namespace math {
 
 /**
- * Return the natural logarithm of the softmax of the specified
- * vector, or of each vector in a container.
- *
+ * Return the natural logarithm of the softmax of the specified vector or 
+ * matrix, or of each vector or matrix in a container. For a matrix, the 
+ * log-softmax is taken over all elements.
+ *  *
  * \f$
  * \log \mbox{softmax}(y)
  * \ = \ y - \log \sum_{k=1}^K \exp(y_k)
@@ -35,17 +36,15 @@ namespace math {
  * \right.
  * \f$
  *
- * @tparam Container type of input: an Eigen vector, `std::vector` of doubles,
- *   or nested container whose scalar type is arithmetic
- * @param x vector or container of vectors to transform
- * @return log softmax of the input, preserving the container structure; an
- *   empty result if any input vector is empty
+ * @tparam Container type of input: an Eigen vector, Eigen matrix, 
+ *   `std::vector` of vectors or matrices, or nested container whose scalar 
+ *   type is arithmetic
+ * @param x vector, matrix, or container to transform.
+ * @return softmax of the input, preserving the container structure; an empty
+ *   result if any input vector or matrix 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>
+          require_container_t<Container>* = nullptr>
 inline auto log_softmax(Container&& x) {
   return make_holder(
       [](auto&& a) {
diff --git a/stan/math/prim/fun/softmax.hpp b/stan/math/prim/fun/softmax.hpp
index 63a51876a60..4271be389ec 100644
--- a/stan/math/prim/fun/softmax.hpp
+++ b/stan/math/prim/fun/softmax.hpp
@@ -11,7 +11,9 @@ namespace stan {
 namespace math {
 
 /**
- * Return the softmax of the specified vector, or of each vector in a container.
+ * Return the softmax of the specified vector or matrix, or of each
+ * vector or matrix in a container. For a matrix, the softmax is
+ * taken over all elements.
  *
  * \f$
  * \mbox{softmax}(y)
@@ -38,17 +40,15 @@ namespace math {
  * \end{array}
  * \f$
  *
- * @tparam Container type of input: an Eigen vector, `std::vector` of doubles,
- *   or nested container whose scalar type is arithmetic
- * @param x vector or container of vectors to transform
+ * @tparam Container type of input: an Eigen vector, Eigen matrix, 
+ *   `std::vector` of vectors or matrices, or nested container whose scalar 
+ *   type is arithmetic
+ * @param x vector, matrix, or container to transform.
  * @return softmax of the input, preserving the container structure; an empty
- *   result if any input vector is empty
+ *   result if any input vector or matrix 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>
+          require_container_t<Container>* = nullptr>
 inline auto softmax(Container&& x) {
   return make_holder(
       [](auto&& a) {
diff --git a/stan/math/rev/fun/log_softmax.hpp b/stan/math/rev/fun/log_softmax.hpp
index 4c77c09e15f..d785f7941c5 100644
--- a/stan/math/rev/fun/log_softmax.hpp
+++ b/stan/math/rev/fun/log_softmax.hpp
@@ -8,15 +8,17 @@
 #include <stan/math/prim/err.hpp>
 #include <stan/math/prim/fun/to_ref.hpp>
 #include <stan/math/prim/fun/log_softmax.hpp>
+#include <stan/math/prim/fun/softmax.hpp>
 #include <stan/math/prim/functor/apply_vector_unary.hpp>
 
 namespace stan {
 namespace math {
 
 /**
- * Return the log softmax of the specified vector or row vector.
+ * Return the log softmax of the specified vector, row vector, or matrix.
  *
- * @tparam T a `var_value` or Eigen vector/row_vector with `var` scalar
+ * @tparam T a `var_value` or Eigen vector, row_vector, or matrix with 
+ *   `var` scalar
  * @param x input
  * @return log softmax of the input, or an empty result if the input is empty
  */
@@ -31,8 +33,9 @@ inline auto log_softmax(T&& x) {
   arena_t<return_t> res = log_softmax(x_arena.val());
   reverse_pass_callback([x_arena, res]() mutable {
     const auto& res_adj = to_ref(res.adj());
+    const auto s = softmax(x_arena.val());
     x_arena.adj().array()
-        += res_adj.array() - res_adj.sum() * res.val().array().exp();
+        += res_adj.array() - res_adj.sum() * s.array();
   });
   return res;
 }
diff --git a/stan/math/rev/fun/log_sum_exp.hpp b/stan/math/rev/fun/log_sum_exp.hpp
index ca6a1b483fb..be9cc20db9e 100644
--- a/stan/math/rev/fun/log_sum_exp.hpp
+++ b/stan/math/rev/fun/log_sum_exp.hpp
@@ -9,6 +9,7 @@
 #include <stan/math/prim/meta.hpp>
 #include <stan/math/prim/fun/constants.hpp>
 #include <stan/math/prim/fun/log_sum_exp.hpp>
+#include <stan/math/prim/fun/softmax.hpp>
 #include <cmath>
 #include <vector>
 
@@ -71,8 +72,7 @@ inline var log_sum_exp(T&& v) {
   auto arena_v_val = to_arena(arena_v.val());
   var res = log_sum_exp(arena_v_val);
   reverse_pass_callback([arena_v, arena_v_val, res]() mutable {
-    arena_v.adj()
-        += res.adj() * (arena_v_val.array().val() - res.val()).exp().matrix();
+    arena_v.adj() += res.adj() * softmax(arena_v_val);
   });
 
   return res;
@@ -87,7 +87,7 @@ inline var log_sum_exp(T&& v) {
 template <typename T, require_var_matrix_t<T>* = nullptr>
 inline var log_sum_exp(const T& x) {
   return make_callback_vari(log_sum_exp(x.val()), [x](const auto& res) mutable {
-    x.adj() += res.adj() * (x.val().array().val() - res.val()).exp().matrix();
+    x.adj() += res.adj() * softmax(x.val());
   });
 }
 
diff --git a/stan/math/rev/fun/softmax.hpp b/stan/math/rev/fun/softmax.hpp
index cd377c9cb1d..074320f3fdc 100644
--- a/stan/math/rev/fun/softmax.hpp
+++ b/stan/math/rev/fun/softmax.hpp
@@ -15,9 +15,10 @@ namespace stan {
 namespace math {
 
 /**
- * Return the softmax of the specified vector or row vector.
+ * Return the softmax of the specified vector, row vector, or matrix.
  *
- * @tparam T a `var_value` or Eigen vector/row_vector with `var` scalar
+ * @tparam T a `var_value` or Eigen vector, row_vector, or matrix with 
+ *   `var` scalar
  * @param x input
  * @return softmax of the input, or an empty result if the input is empty
  */
@@ -31,17 +32,18 @@ inline auto softmax(T&& x) {
       = return_var_matrix_t<plain_type_t<decltype(x_arena.val())>, T>;
   arena_t<return_t> res = softmax(x_arena.val());
   reverse_pass_callback([x_arena, res]() mutable {
+    const auto dot_sd = (res.val().array() * res.adj().array()).sum();
     x_arena.adj().array()
-        += res.val().array() * (res.adj().array() - res.val().dot(res.adj()));
+        += res.val().array() * (res.adj().array() - dot_sd);
   });
   return res;
 }
 
 /**
- * Return the softmax of each vector in an array.
+ * Return the softmax of each vector or matrix in an array.
  *
  * @tparam T `std::vector` whose scalar type is `var`
- * @param x array of vectors to transform
+ * @param x array of vectors or matrices to transform
  * @return array of softmax results
  */
 template <typename T, require_std_vector_st<is_var, T>* = nullptr>

From 47e34d1a37197ca2dfe135759ff9e9776ad3ce59 Mon Sep 17 00:00:00 2001
From: Jouni Helske <jvhels@utu.fi>
Date: Thu, 22 Oct 2026 12:22:11 +0300
Subject: [PATCH 6/7] fixes based on review, helper for sofmax and lse

---
 .gitignore                                    |  1 +
 stan/math/fwd/fun/log_softmax.hpp             | 35 ++++------
 stan/math/fwd/fun/log_sum_exp.hpp             |  8 +--
 stan/math/fwd/fun/softmax.hpp                 | 25 +++----
 stan/math/prim/fun.hpp                        |  1 +
 stan/math/prim/fun/log_softmax.hpp            | 14 ++--
 stan/math/prim/fun/softmax.hpp                |  6 +-
 stan/math/prim/fun/softmax_log_sum_exp.hpp    | 69 +++++++++++++++++++
 stan/math/rev/fun/log_softmax.hpp             |  2 +-
 stan/math/rev/fun/log_sum_exp.hpp             | 20 +++---
 stan/math/rev/fun/softmax.hpp                 |  2 +-
 test/unit/math/fwd/fun/log_softmax_test.cpp   | 57 ++++-----------
 test/unit/math/fwd/fun/log_sum_exp_test.cpp   | 10 ---
 .../prim/fun/softmax_log_sum_exp_test.cpp     | 59 ++++++++++++++++
 14 files changed, 191 insertions(+), 118 deletions(-)
 create mode 100644 stan/math/prim/fun/softmax_log_sum_exp.hpp
 create mode 100644 test/unit/math/prim/fun/softmax_log_sum_exp_test.cpp

diff --git a/.gitignore b/.gitignore
index c36591f03d8..63af199a251 100644
--- a/.gitignore
+++ b/.gitignore
@@ -106,3 +106,4 @@ lib/tbb/**
 
 # gdb debug history
 .gdb_history
+.positai
diff --git a/stan/math/fwd/fun/log_softmax.hpp b/stan/math/fwd/fun/log_softmax.hpp
index 24c96922b64..8b96e9a231c 100644
--- a/stan/math/fwd/fun/log_softmax.hpp
+++ b/stan/math/fwd/fun/log_softmax.hpp
@@ -4,9 +4,9 @@
 #include <stan/math/prim/fun/Eigen.hpp>
 #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/softmax_log_sum_exp.hpp>
 #include <stan/math/prim/fun/to_ref.hpp>
 #include <stan/math/prim/functor/apply_vector_unary.hpp>
 
@@ -14,7 +14,7 @@ namespace stan {
 namespace math {
 
 /**
- * Return the log softmax of each vector or matrix in a container of `fvar` 
+ * Return the log softmax of each vector or matrix in a container of `fvar`
  * values.
  *
  * @tparam T `std::vector` whose scalar type is `fvar`
@@ -33,31 +33,22 @@ inline auto log_softmax(T&& x) {
  *
  * @tparam Mat Eigen vector or matrix with `fvar` scalar
  * @param x vector or matrix to transform
- * @return log softmax of the vector or matrix, or an empty result if the 
+ * @return log softmax of the vector or matrix, or an empty result if the
  * input is empty
  */
 template <typename Mat, require_eigen_vt<is_fvar, Mat>* = nullptr>
-inline auto log_softmax(Mat&& x) {
-  using mat = std::decay_t<Mat>;
-  constexpr int Rows = mat::RowsAtCompileTime;
-  constexpr int Cols = mat::ColsAtCompileTime;
-  using T = typename value_type_t<mat>::Scalar;
-  decltype(auto) x_ref = to_ref(std::forward<Mat>(x));
-  if (x_ref.size() == 0) {
-    return Eigen::Matrix<fvar<T>, Rows, Cols>{};
+inline plain_type_t<Mat> log_softmax(Mat&& x) {
+  if (x.size() == 0) {
+    return {};
   }
-  const auto x_val = value_of(x_ref);
-  const auto lse = log_sum_exp(x_val);
-  const auto s = softmax(x_val);
-  const auto d_in = x_ref.d();
-  const auto dot_sd = (s.array() * d_in.array()).sum();
-  
-  Eigen::Matrix<fvar<T>, Rows, Cols> result(x_ref.rows(), x_ref.cols());
-  result.val() = (x_val.array() - lse).matrix();
-  result.d() = (d_in.array() - dot_sd).matrix();
-  return result;
+  decltype(auto) x_ref = to_ref(std::forward<Mat>(x));
+  const auto x_val = value_of(x_ref).eval();
+  const auto d_in = x_ref.d().eval();
+  const auto softmax_lse = internal::softmax_log_sum_exp(x_val);
+  return to_fvar(x_val.array() - softmax_lse.second,
+                 d_in.array() - softmax_lse.first.cwiseProduct(d_in).sum());
 }
 
 }  // namespace math
 }  // namespace stan
-#endif
\ No newline at end of file
+#endif
diff --git a/stan/math/fwd/fun/log_sum_exp.hpp b/stan/math/fwd/fun/log_sum_exp.hpp
index 399b13e877e..88498ea2c70 100644
--- a/stan/math/fwd/fun/log_sum_exp.hpp
+++ b/stan/math/fwd/fun/log_sum_exp.hpp
@@ -8,7 +8,7 @@
 #include <stan/math/prim/fun/constants.hpp>
 #include <stan/math/prim/fun/to_ref.hpp>
 #include <stan/math/prim/fun/log_sum_exp.hpp>
-#include <stan/math/prim/fun/softmax.hpp>
+#include <stan/math/prim/fun/softmax_log_sum_exp.hpp>
 #include <cmath>
 #include <vector>
 
@@ -57,11 +57,9 @@ inline auto log_sum_exp(T&& x) {
         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();
-        
-        const auto probs = softmax(vals);
+        const auto softmax_lse = internal::softmax_log_sum_exp(vals);
         return fvar<T_fvar_inner>(
-          log_sum_exp(vals),
-          v.d().cwiseProduct(probs).sum());
+            softmax_lse.second, v.d().cwiseProduct(softmax_lse.first).sum());
       });
 }
 
diff --git a/stan/math/fwd/fun/softmax.hpp b/stan/math/fwd/fun/softmax.hpp
index c77cb35bf33..34e35adbc80 100644
--- a/stan/math/fwd/fun/softmax.hpp
+++ b/stan/math/fwd/fun/softmax.hpp
@@ -31,26 +31,19 @@ inline auto softmax(T&& x) {
  *
  * @tparam Mat Eigen vector or matrix with `fvar` scalar
  * @param x vector or matrix to transform
- * @return softmax of the vector, matrix, or an empty result if the input is 
+ * @return softmax of the vector, matrix, or an empty result if the input is
  *   empty
  */
 template <typename Mat, require_eigen_vt<is_fvar, Mat>* = nullptr>
-inline auto softmax(Mat&& x) {
-  using mat = std::decay_t<Mat>;
-  constexpr int Rows = mat::RowsAtCompileTime;
-  constexpr int Cols = mat::ColsAtCompileTime;
-  using T = typename value_type_t<mat>::Scalar;
-  decltype(auto) x_ref = to_ref(std::forward<Mat>(x));
-  if (x_ref.size() == 0) {
-    return Eigen::Matrix<fvar<T>, Rows, Cols>{};
+inline plain_type_t<Mat> softmax(Mat&& x) {
+  if (x.size() == 0) {
+    return {};
   }
-  const auto s = softmax(value_of(x_ref));
-  const auto d_in = x_ref.d();
-  const auto dot_sd = (s.array() * d_in.array()).sum();
-  Eigen::Matrix<fvar<T>, Rows, Cols> result(x_ref.rows(), x_ref.cols());
-  result.val() = s;
-  result.d() = (s.array() * (d_in.array() - dot_sd)).matrix();
-  return result;
+  decltype(auto) x_ref = to_ref(std::forward<Mat>(x));
+  const auto s = softmax(value_of(x_ref)).eval();
+  const auto d_in = x_ref.d().eval();
+
+  return to_fvar(s, s.array() * (d_in.array() - s.cwiseProduct(d_in).sum()));
 }
 
 }  // namespace math
diff --git a/stan/math/prim/fun.hpp b/stan/math/prim/fun.hpp
index 91b039639da..31dc5370689 100644
--- a/stan/math/prim/fun.hpp
+++ b/stan/math/prim/fun.hpp
@@ -295,6 +295,7 @@
 #include <stan/math/prim/fun/size_mvt.hpp>
 #include <stan/math/prim/fun/size_zero.hpp>
 #include <stan/math/prim/fun/softmax.hpp>
+#include <stan/math/prim/fun/softmax_log_sum_exp.hpp>
 #include <stan/math/prim/fun/sort_asc.hpp>
 #include <stan/math/prim/fun/sort_desc.hpp>
 #include <stan/math/prim/fun/sort_indices.hpp>
diff --git a/stan/math/prim/fun/log_softmax.hpp b/stan/math/prim/fun/log_softmax.hpp
index 51665b47adb..c3760135df7 100644
--- a/stan/math/prim/fun/log_softmax.hpp
+++ b/stan/math/prim/fun/log_softmax.hpp
@@ -12,10 +12,10 @@ namespace stan {
 namespace math {
 
 /**
- * Return the natural logarithm of the softmax of the specified vector or 
- * matrix, or of each vector or matrix in a container. For a matrix, the 
+ * Return the natural logarithm of the softmax of the specified vector or
+ * matrix, or of each vector or matrix in a container. For a matrix, the
  * log-softmax is taken over all elements.
- *  *
+ *
  * \f$
  * \log \mbox{softmax}(y)
  * \ = \ y - \log \sum_{k=1}^K \exp(y_k)
@@ -36,12 +36,12 @@ namespace math {
  * \right.
  * \f$
  *
- * @tparam Container type of input: an Eigen vector, Eigen matrix, 
- *   `std::vector` of vectors or matrices, or nested container whose scalar 
+ * @tparam Container type of input: an Eigen vector, Eigen matrix,
+ *   `std::vector` of vectors or matrices, or nested container whose scalar
  *   type is arithmetic
  * @param x vector, matrix, or container to transform.
- * @return softmax of the input, preserving the container structure; an empty
- *   result if any input vector or matrix is empty.
+ * @return log softmax of the input, preserving the container structure; an
+ *   empty result if any input vector or matrix is empty.
  */
 template <typename Container, require_st_arithmetic<Container>* = nullptr,
           require_container_t<Container>* = nullptr>
diff --git a/stan/math/prim/fun/softmax.hpp b/stan/math/prim/fun/softmax.hpp
index 4271be389ec..6cac82507be 100644
--- a/stan/math/prim/fun/softmax.hpp
+++ b/stan/math/prim/fun/softmax.hpp
@@ -40,8 +40,8 @@ namespace math {
  * \end{array}
  * \f$
  *
- * @tparam Container type of input: an Eigen vector, Eigen matrix, 
- *   `std::vector` of vectors or matrices, or nested container whose scalar 
+ * @tparam Container type of input: an Eigen vector, Eigen matrix,
+ *   `std::vector` of vectors or matrices, or nested container whose scalar
  *   type is arithmetic
  * @param x vector, matrix, or container to transform.
  * @return softmax of the input, preserving the container structure; an empty
@@ -58,7 +58,7 @@ inline auto softmax(Container&& x) {
               if (v.size() == 0) {
                 return v;
               }
-              const auto theta = (v.array() - v.maxCoeff()).exp();
+              const auto theta = (v.array() - v.maxCoeff()).exp().eval();
               return (theta / theta.sum()).matrix();
             });
       },
diff --git a/stan/math/prim/fun/softmax_log_sum_exp.hpp b/stan/math/prim/fun/softmax_log_sum_exp.hpp
new file mode 100644
index 00000000000..d4ac4f094d9
--- /dev/null
+++ b/stan/math/prim/fun/softmax_log_sum_exp.hpp
@@ -0,0 +1,69 @@
+#ifndef STAN_MATH_PRIM_FUN_SOFTMAX_LOG_SUM_EXP_HPP
+#define STAN_MATH_PRIM_FUN_SOFTMAX_LOG_SUM_EXP_HPP
+
+#include <stan/math/prim/meta.hpp>
+#include <stan/math/prim/fun/constants.hpp>
+#include <stan/math/prim/fun/Eigen.hpp>
+#include <stan/math/prim/fun/log_sum_exp.hpp>
+#include <stan/math/prim/fun/softmax.hpp>
+#include <stan/math/prim/fun/to_ref.hpp>
+#include <cmath>
+#include <utility>
+
+namespace stan {
+namespace math {
+namespace internal {
+
+/**
+ * Return the softmax and the log sum of exponentials of the specified
+ * vector or matrix, using a single evaluation of the exponential per
+ * element. For a matrix, both are taken over all elements.
+ *
+ * With \f$m = \max(x)\f$ and \f$\theta = \exp(x - m)\f$,
+ *
+ * \f$\mbox{softmax}(x) = \theta / \sum_{n=1}^N \theta_n\f$ and
+ * \f$\mbox{log\_sum\_exp}(x) = m + \log \sum_{n=1}^N \theta_n\f$.
+ *
+ * The results match `softmax(x)` and `log_sum_exp(x)`, including for
+ * non-finite input: if \f$m\f$ is not finite, the log sum of exponentials
+ * is \f$m\f$ and the softmax is NaN.
+ *
+ * @tparam T type of Eigen vector or matrix
+ * @param x Eigen vector or matrix
+ * @return pair of the softmax of `x` and the log sum of exponentials of
+ *   `x`. For empty input, returns an empty softmax and negative infinity
+ */
+template <typename T, require_eigen_vt<std::is_arithmetic, T>* = nullptr>
+inline std::pair<plain_type_t<T>, double> softmax_log_sum_exp(const T& x) {
+  if (x.size() == 0) {
+    return {x, NEGATIVE_INFTY};
+  }
+  const auto& x_ref = to_ref(x);
+  const double max = x_ref.maxCoeff();
+  const auto theta = (x_ref.array() - max).exp().eval();
+  const double theta_sum = theta.sum();
+  const double lse = std::isfinite(max) ? max + std::log(theta_sum) : max;
+  return {plain_type_t<T>((theta / theta_sum).matrix()), lse};
+}
+
+/**
+ * Return the softmax and the log sum of exponentials of the specified
+ * vector or matrix of autodiff values. This happens with nested autodiff,
+ * e.g. the values of an `fvar<var>` vector, and falls back to the
+ * autodiff specializations of `softmax()` and `log_sum_exp()`.
+ *
+ * @tparam T type of Eigen vector or matrix with autodiff scalar
+ * @param x vector or matrix
+ * @return pair of the softmax of `x` and the log sum of exponentials of `x`
+ */
+template <typename T, require_eigen_vt<is_autodiff, T>* = nullptr>
+inline std::pair<plain_type_t<T>, value_type_t<T>> softmax_log_sum_exp(
+    const T& x) {
+  return {plain_type_t<T>(softmax(x)), log_sum_exp(x)};
+}
+
+}  // namespace internal
+}  // namespace math
+}  // namespace stan
+
+#endif
diff --git a/stan/math/rev/fun/log_softmax.hpp b/stan/math/rev/fun/log_softmax.hpp
index d785f7941c5..ccde47b630b 100644
--- a/stan/math/rev/fun/log_softmax.hpp
+++ b/stan/math/rev/fun/log_softmax.hpp
@@ -17,7 +17,7 @@ namespace math {
 /**
  * Return the log softmax of the specified vector, row vector, or matrix.
  *
- * @tparam T a `var_value` or Eigen vector, row_vector, or matrix with 
+ * @tparam T a `var_value` or Eigen vector, row_vector, or matrix with
  *   `var` scalar
  * @param x input
  * @return log softmax of the input, or an empty result if the input is empty
diff --git a/stan/math/rev/fun/log_sum_exp.hpp b/stan/math/rev/fun/log_sum_exp.hpp
index be9cc20db9e..fc7f770c5d1 100644
--- a/stan/math/rev/fun/log_sum_exp.hpp
+++ b/stan/math/rev/fun/log_sum_exp.hpp
@@ -9,7 +9,7 @@
 #include <stan/math/prim/meta.hpp>
 #include <stan/math/prim/fun/constants.hpp>
 #include <stan/math/prim/fun/log_sum_exp.hpp>
-#include <stan/math/prim/fun/softmax.hpp>
+#include <stan/math/prim/fun/softmax_log_sum_exp.hpp>
 #include <cmath>
 #include <vector>
 
@@ -69,10 +69,11 @@ template <typename T, require_eigen_st<is_var, T>* = nullptr,
           require_not_var_matrix_t<T>* = nullptr>
 inline var log_sum_exp(T&& v) {
   arena_t<T> arena_v = std::forward<T>(v);
-  auto arena_v_val = to_arena(arena_v.val());
-  var res = log_sum_exp(arena_v_val);
-  reverse_pass_callback([arena_v, arena_v_val, res]() mutable {
-    arena_v.adj() += res.adj() * softmax(arena_v_val);
+  const auto softmax_lse = internal::softmax_log_sum_exp(arena_v.val());
+  auto arena_softmax = to_arena(softmax_lse.first);
+  var res = softmax_lse.second;
+  reverse_pass_callback([arena_v, arena_softmax, res]() mutable {
+    arena_v.adj() += res.adj() * arena_softmax;
   });
 
   return res;
@@ -86,9 +87,12 @@ inline var log_sum_exp(T&& v) {
  */
 template <typename T, require_var_matrix_t<T>* = nullptr>
 inline var log_sum_exp(const T& x) {
-  return make_callback_vari(log_sum_exp(x.val()), [x](const auto& res) mutable {
-    x.adj() += res.adj() * softmax(x.val());
-  });
+  const auto softmax_lse = internal::softmax_log_sum_exp(x.val());
+  auto arena_softmax = to_arena(softmax_lse.first);
+  return make_callback_vari(softmax_lse.second,
+                            [x, arena_softmax](const auto& res) mutable {
+                              x.adj() += res.adj() * arena_softmax;
+                            });
 }
 
 /**
diff --git a/stan/math/rev/fun/softmax.hpp b/stan/math/rev/fun/softmax.hpp
index 074320f3fdc..4c908094408 100644
--- a/stan/math/rev/fun/softmax.hpp
+++ b/stan/math/rev/fun/softmax.hpp
@@ -17,7 +17,7 @@ namespace math {
 /**
  * Return the softmax of the specified vector, row vector, or matrix.
  *
- * @tparam T a `var_value` or Eigen vector, row_vector, or matrix with 
+ * @tparam T a `var_value` or Eigen vector, row_vector, or matrix with
  *   `var` scalar
  * @param x input
  * @return softmax of the input, or an empty result if the input is empty
diff --git a/test/unit/math/fwd/fun/log_softmax_test.cpp b/test/unit/math/fwd/fun/log_softmax_test.cpp
index 9055b896af2..59ebc16420d 100644
--- a/test/unit/math/fwd/fun/log_softmax_test.cpp
+++ b/test/unit/math/fwd/fun/log_softmax_test.cpp
@@ -1,48 +1,26 @@
 #include <stan/math/fwd.hpp>
+#include <test/unit/util.hpp>
 #include <gtest/gtest.h>
 
 #include <cmath>
 
 TEST(AgradFwd, log_softmax_large_dynamic_range) {
   using stan::math::fvar;
+  using stan::math::vector_fd;
   using stan::math::log_softmax;
 
-  Eigen::Matrix<fvar<double>, Eigen::Dynamic, 1> x(3);
+  Eigen::VectorXd val(3);
+  val << 0, -100, -1000;
 
-  x(0) = fvar<double>(0.0, 1.0);
-  x(1) = fvar<double>(-100.0, 2.0);
-  x(2) = fvar<double>(-1000.0, 3.0);
+  Eigen::VectorXd d(3);
+  d << 1, 2, 3;
+
+  vector_fd x = stan::math::to_fvar(val, d);
 
   auto y = log_softmax(x);
 
-  /*
-   * log_softmax(x) = x - log_sum_exp(x).
-   *
-   * exp(-1000) underflows to zero in double precision. Therefore:
-   *
-   *   log(softmax(x)[2])
-   *
-   * gives -inf while
-   *
-   *   x[2] - log_sum_exp(x)
-   *
-   * correctly retains the log-probability.
-   */
-
-  EXPECT_TRUE(std::isfinite(y.val()(0)));
-  EXPECT_TRUE(std::isfinite(y.val()(1)));
-  EXPECT_TRUE(std::isfinite(y.val()(2)));
-
-  EXPECT_DOUBLE_EQ(0.0, y.val()(0));
-  EXPECT_DOUBLE_EQ(-100.0, y.val()(1));
-  EXPECT_DOUBLE_EQ(-1000.0, y.val()(2));
-
-  const double p1 = std::exp(-100.0);
-  const double dot_sd = 1.0 + 2.0 * p1;
-
-  EXPECT_NEAR(1.0 - dot_sd, y(0).d(), 1e-12);
-  EXPECT_NEAR(2.0 - dot_sd, y(1).d(), 1e-12);
-  EXPECT_NEAR(3.0 - dot_sd, y(2).d(), 1e-12);
+  EXPECT_MATRIX_EQ(y.val(), val);
+  EXPECT_MATRIX_EQ(y.d(), d.array() - (1.0 + 2.0 * std::exp(-100.0)));
 }
 
 TEST(AgradFwd, log_softmax_many_small_probabilities) {
@@ -79,6 +57,7 @@ TEST(AgradFwd, log_softmax_many_small_probabilities) {
   EXPECT_NEAR(1.0 - tangent, y(0).d(), 1e-12);
   EXPECT_NEAR(2.0 - tangent, y(1).d(), 1e-12);
 }
+
 TEST(AgradFwd, log_softmax_positive_infinity_consistent_with_prim) {
   using stan::math::fvar;
   using stan::math::log_softmax;
@@ -93,20 +72,8 @@ TEST(AgradFwd, log_softmax_positive_infinity_consistent_with_prim) {
 
   auto y = log_softmax(x);
 
-  // The value is defined as
-  //
-  //   x - log_sum_exp(x).
-  //
-  // Since log_sum_exp([1, 2, 3, +Inf]) = +Inf:
-  //
-  //   1 - Inf = -Inf
-  //   2 - Inf = -Inf
-  //   3 - Inf = -Inf
-  //   Inf - Inf = NaN
-  //
-
   EXPECT_EQ(-inf, y(0).val());
   EXPECT_EQ(-inf, y(1).val());
   EXPECT_EQ(-inf, y(2).val());
   EXPECT_TRUE(std::isnan(y(3).val()));
-}
\ No newline at end of file
+}
diff --git a/test/unit/math/fwd/fun/log_sum_exp_test.cpp b/test/unit/math/fwd/fun/log_sum_exp_test.cpp
index e6286602bb5..3da59876486 100644
--- a/test/unit/math/fwd/fun/log_sum_exp_test.cpp
+++ b/test/unit/math/fwd/fun/log_sum_exp_test.cpp
@@ -13,22 +13,12 @@ TEST(AgradFwd, log_sum_exp_derivative_does_not_overflow) {
   x(2) = fvar<double>(999.0, 3.0);
 
   auto y = log_sum_exp(x);
-  // log_sum_exp([1000, 1001, 999])
-  //
-  // = 1001 + log(exp(-1) + 1 + exp(-2)).
   const double expected_value
       = 1001.0 + std::log(std::exp(-1.0) + 1.0 + std::exp(-2.0));
 
   EXPECT_NEAR(expected_value, y.val(), 1e-12);
   EXPECT_TRUE(std::isfinite(y.val()));
 
-  // The derivative is
-  //
-  // softmax(x) dot x.d().
-  //
-  // The old implementation formed exp(x) directly in the derivative,
-  // which gives inf / inf for these values.
-
   const double denom = std::exp(-1.0) + 1.0 + std::exp(-2.0);
 
   const double p0 = std::exp(-1.0) / denom;
diff --git a/test/unit/math/prim/fun/softmax_log_sum_exp_test.cpp b/test/unit/math/prim/fun/softmax_log_sum_exp_test.cpp
new file mode 100644
index 00000000000..6b08d4e9828
--- /dev/null
+++ b/test/unit/math/prim/fun/softmax_log_sum_exp_test.cpp
@@ -0,0 +1,59 @@
+#include <stan/math/prim.hpp>
+#include <test/unit/util.hpp>
+#include <gtest/gtest.h>
+#include <cmath>
+#include <limits>
+
+TEST(MathFunctions, softmax_log_sum_exp_matches_separate_calls) {
+  using stan::math::internal::softmax_log_sum_exp;
+  using stan::math::log_sum_exp;
+  using stan::math::softmax;
+
+  Eigen::VectorXd v(4);
+  v << 629.7901581243797, 31.52411463, 608.19720553, 120.94829574;
+  auto res = softmax_log_sum_exp(v);
+  EXPECT_MATRIX_EQ(softmax(v), res.first);
+  EXPECT_FLOAT_EQ(log_sum_exp(v), res.second);
+
+  Eigen::RowVectorXd rv(3);
+  rv << -1.0, 1.0, 10.0;
+  auto res_rv = softmax_log_sum_exp(rv);
+  EXPECT_MATRIX_EQ(softmax(rv), res_rv.first);
+  EXPECT_FLOAT_EQ(log_sum_exp(rv), res_rv.second);
+
+  Eigen::MatrixXd m(2, 2);
+  m << 0.0, -100.0, -1000.0, 3.0;
+  auto res_m = softmax_log_sum_exp(m);
+  EXPECT_MATRIX_EQ(softmax(m), res_m.first);
+  EXPECT_FLOAT_EQ(log_sum_exp(m), res_m.second);
+}
+
+TEST(MathFunctions, softmax_log_sum_exp_empty) {
+  using stan::math::internal::softmax_log_sum_exp;
+
+  Eigen::VectorXd v(0);
+  auto res = softmax_log_sum_exp(v);
+  EXPECT_EQ(0, res.first.size());
+  EXPECT_EQ(stan::math::NEGATIVE_INFTY, res.second);
+}
+
+TEST(MathFunctions, softmax_log_sum_exp_non_finite) {
+  using stan::math::internal::softmax_log_sum_exp;
+  constexpr double inf = std::numeric_limits<double>::infinity();
+
+  Eigen::VectorXd v(3);
+  v << 1.0, -inf, 2.0;
+  auto res = softmax_log_sum_exp(v);
+  EXPECT_FLOAT_EQ(2.0 + std::log(std::exp(-1.0) + 1.0), res.second);
+  EXPECT_EQ(0.0, res.first(1));
+
+  v << 1.0, inf, 2.0;
+  res = softmax_log_sum_exp(v);
+  EXPECT_EQ(inf, res.second);
+  EXPECT_TRUE(res.first.array().isNaN().all());
+
+  v << -inf, -inf, -inf;
+  res = softmax_log_sum_exp(v);
+  EXPECT_EQ(-inf, res.second);
+  EXPECT_TRUE(res.first.array().isNaN().all());
+}

From 6623432dc2fa564fc0d0dc63405342413dd5072a Mon Sep 17 00:00:00 2001
From: Stan Jenkins <mc.stanislaw@gmail.com>
Date: Thu, 22 Oct 2026 06:15:26 -0400
Subject: [PATCH 7/7] [Jenkins] auto-formatting by clang-format version
 10.0.0-4ubuntu1

---
 stan/math/fwd/fun/log_sum_exp.hpp                    | 4 ++--
 stan/math/rev/fun/log_softmax.hpp                    | 3 +--
 stan/math/rev/fun/softmax.hpp                        | 3 +--
 test/unit/math/fwd/fun/log_softmax_test.cpp          | 2 +-
 test/unit/math/prim/fun/softmax_log_sum_exp_test.cpp | 2 +-
 5 files changed, 6 insertions(+), 8 deletions(-)

diff --git a/stan/math/fwd/fun/log_sum_exp.hpp b/stan/math/fwd/fun/log_sum_exp.hpp
index 88498ea2c70..1e5f01a2259 100644
--- a/stan/math/fwd/fun/log_sum_exp.hpp
+++ b/stan/math/fwd/fun/log_sum_exp.hpp
@@ -58,8 +58,8 @@ inline auto log_sum_exp(T&& x) {
         using mat_type = Eigen::Matrix<T_fvar_inner, -1, -1>;
         mat_type vals = v.val();
         const auto softmax_lse = internal::softmax_log_sum_exp(vals);
-        return fvar<T_fvar_inner>(
-            softmax_lse.second, v.d().cwiseProduct(softmax_lse.first).sum());
+        return fvar<T_fvar_inner>(softmax_lse.second,
+                                  v.d().cwiseProduct(softmax_lse.first).sum());
       });
 }
 
diff --git a/stan/math/rev/fun/log_softmax.hpp b/stan/math/rev/fun/log_softmax.hpp
index ccde47b630b..845d03aa147 100644
--- a/stan/math/rev/fun/log_softmax.hpp
+++ b/stan/math/rev/fun/log_softmax.hpp
@@ -34,8 +34,7 @@ inline auto log_softmax(T&& x) {
   reverse_pass_callback([x_arena, res]() mutable {
     const auto& res_adj = to_ref(res.adj());
     const auto s = softmax(x_arena.val());
-    x_arena.adj().array()
-        += res_adj.array() - res_adj.sum() * s.array();
+    x_arena.adj().array() += res_adj.array() - res_adj.sum() * s.array();
   });
   return res;
 }
diff --git a/stan/math/rev/fun/softmax.hpp b/stan/math/rev/fun/softmax.hpp
index 4c908094408..d3d46665774 100644
--- a/stan/math/rev/fun/softmax.hpp
+++ b/stan/math/rev/fun/softmax.hpp
@@ -33,8 +33,7 @@ inline auto softmax(T&& x) {
   arena_t<return_t> res = softmax(x_arena.val());
   reverse_pass_callback([x_arena, res]() mutable {
     const auto dot_sd = (res.val().array() * res.adj().array()).sum();
-    x_arena.adj().array()
-        += res.val().array() * (res.adj().array() - dot_sd);
+    x_arena.adj().array() += res.val().array() * (res.adj().array() - dot_sd);
   });
   return res;
 }
diff --git a/test/unit/math/fwd/fun/log_softmax_test.cpp b/test/unit/math/fwd/fun/log_softmax_test.cpp
index 59ebc16420d..dd713ed1641 100644
--- a/test/unit/math/fwd/fun/log_softmax_test.cpp
+++ b/test/unit/math/fwd/fun/log_softmax_test.cpp
@@ -6,8 +6,8 @@
 
 TEST(AgradFwd, log_softmax_large_dynamic_range) {
   using stan::math::fvar;
-  using stan::math::vector_fd;
   using stan::math::log_softmax;
+  using stan::math::vector_fd;
 
   Eigen::VectorXd val(3);
   val << 0, -100, -1000;
diff --git a/test/unit/math/prim/fun/softmax_log_sum_exp_test.cpp b/test/unit/math/prim/fun/softmax_log_sum_exp_test.cpp
index 6b08d4e9828..4a17014a3a6 100644
--- a/test/unit/math/prim/fun/softmax_log_sum_exp_test.cpp
+++ b/test/unit/math/prim/fun/softmax_log_sum_exp_test.cpp
@@ -5,9 +5,9 @@
 #include <limits>
 
 TEST(MathFunctions, softmax_log_sum_exp_matches_separate_calls) {
-  using stan::math::internal::softmax_log_sum_exp;
   using stan::math::log_sum_exp;
   using stan::math::softmax;
+  using stan::math::internal::softmax_log_sum_exp;
 
   Eigen::VectorXd v(4);
   v << 629.7901581243797, 31.52411463, 608.19720553, 120.94829574;