From 9f783cd807f7d80bf36a845a237d72c54df2ab5e Mon Sep 17 00:00:00 2001
From: Vaibhav Pathak <pathakvaibhav@protonmail.com>
Date: Mon, 10 Apr 2023 06:35:33 +0530
Subject: [PATCH 1/2] Fix Backward and fix test

---
 src/mlpack/methods/ann/layer/log_softmax_impl.hpp |  2 +-
 src/mlpack/tests/ann/layer/log_softmax.cpp        | 11 +++++------
 2 files changed, 6 insertions(+), 7 deletions(-)

diff --git a/src/mlpack/methods/ann/layer/log_softmax_impl.hpp b/src/mlpack/methods/ann/layer/log_softmax_impl.hpp
index e9516b6fb3..7efc4005b4 100644
--- a/src/mlpack/methods/ann/layer/log_softmax_impl.hpp
+++ b/src/mlpack/methods/ann/layer/log_softmax_impl.hpp
@@ -103,7 +103,7 @@ void LogSoftMaxType<MatType>::Backward(
     const MatType& gy,
     MatType& g)
 {
-  g = arma::exp(input) + gy;
+  g = gy - arma::exp(input) % arma::repmat(arma::sum(gy), input.n_rows, 1);
 }
 
 } // namespace mlpack
diff --git a/src/mlpack/tests/ann/layer/log_softmax.cpp b/src/mlpack/tests/ann/layer/log_softmax.cpp
index d0241d66fd..9c8662ba11 100644
--- a/src/mlpack/tests/ann/layer/log_softmax.cpp
+++ b/src/mlpack/tests/ann/layer/log_softmax.cpp
@@ -29,16 +29,15 @@ TEST_CASE("SimpleLogSoftmaxLayerTest", "[ANNLayerTest]")
   LogSoftMax module;
 
   // Test the Forward function.
-  input = arma::mat("0.5; 0.5");
+  input = arma::mat("-0.6871; 0.7898; 0.2011; 0.0949; -0.0550");
   module.Forward(input, output);
-  REQUIRE(arma::accu(arma::abs(arma::mat("-0.6931; -0.6931") - output)) ==
+  REQUIRE(arma::accu(arma::abs(arma::mat("-2.4746; -0.9977; -1.5864; -1.6926; -1.8425") - output)) ==
       Approx(0.0).margin(1e-3));
 
   // Test the Backward function.
-  error = arma::zeros(input.n_rows, input.n_cols);
+  error = arma::ones(input.n_rows, input.n_cols);
   // Assume LogSoftmax layer is always associated with NLL output layer.
-  error(1, 0) = -1;
-  module.Backward(input, error, delta);
-  REQUIRE(arma::accu(arma::abs(arma::mat("1.6487; 0.6487") - delta)) ==
+  module.Backward(output, error, delta);
+  REQUIRE(arma::accu(arma::abs(arma::mat("0.5790; -0.8435; -0.0233; 0.0798; 0.2079") - delta)) ==
       Approx(0.0).margin(1e-3));
 }

From 8128277f438ec9e2ed019d9fde25e689ae6d878f Mon Sep 17 00:00:00 2001
From: Vaibhav Pathak <pathakvaibhav@protonmail.com>
Date: Sat, 29 Apr 2023 03:19:08 +0530
Subject: [PATCH 2/2] Add JacobianTest

---
 src/mlpack/tests/ann/layer/log_softmax.cpp | 21 +++++++++++++++++++++
 1 file changed, 21 insertions(+)

diff --git a/src/mlpack/tests/ann/layer/log_softmax.cpp b/src/mlpack/tests/ann/layer/log_softmax.cpp
index 9c8662ba11..1f5c3acc72 100644
--- a/src/mlpack/tests/ann/layer/log_softmax.cpp
+++ b/src/mlpack/tests/ann/layer/log_softmax.cpp
@@ -41,3 +41,24 @@ TEST_CASE("SimpleLogSoftmaxLayerTest", "[ANNLayerTest]")
   REQUIRE(arma::accu(arma::abs(arma::mat("0.5790; -0.8435; -0.0233; 0.0798; 0.2079") - delta)) ==
       Approx(0.0).margin(1e-3));
 }
+
+/**
+ * JacobianTest for LogSoftMax layer
+ */
+TEST_CASE("JacobianLogSoftMaxLayerTest", "[ANNLayerTest]")
+{
+  for (size_t i = 0; i < 5; ++i)
+  {
+    const size_t elems = arma::randi(arma::distr_param(2, 1000));
+
+    arma::mat input(elems, 1);
+
+    LogSoftMax module;
+    module.InputDimensions() = { elems };
+    module.ComputeOutputDimensions();
+
+    double error = JacobianTest(module, input);
+    REQUIRE(error <= 1e-5);
+  }
+}
+