From 600d207aa25ecece5133f65f77232e7e03b4a18d Mon Sep 17 00:00:00 2001
From: Anwaar <shahanwaar.khalid8@gmail.com>
Date: Mon, 7 Feb 2022 18:10:55 +0530
Subject: [PATCH] Poisson NLL.

---
 .../ann/loss_functions/poisson_nll_loss.hpp   | 25 +++++++++++--------
 .../loss_functions/poisson_nll_loss_impl.hpp  | 16 +++++++-----
 2 files changed, 24 insertions(+), 17 deletions(-)

diff --git a/src/mlpack/methods/ann/loss_functions/poisson_nll_loss.hpp b/src/mlpack/methods/ann/loss_functions/poisson_nll_loss.hpp
index 31e1cb5620..aa0098761a 100644
--- a/src/mlpack/methods/ann/loss_functions/poisson_nll_loss.hpp
+++ b/src/mlpack/methods/ann/loss_functions/poisson_nll_loss.hpp
@@ -45,12 +45,16 @@ class PoissonNLLLoss
    * @param full Boolean value that determines whether to include Stirling's
    *        approximation term.
    * @param eps A small value to prevent 0 in denominators and logarithms.
-   * @param mean When true, mean loss is computed otherwise total loss.
+   * @param reduction Specifies the reduction to apply to the output. If false,
+   *                  'mean' reduction is used, where sum of the output will be
+   *                  divided by the number of elements in the output. If true,
+   *                  'sum' reduction is used and the output will be summed. It
+   *                  is set to true by default.
    */
   PoissonNLLLoss(const bool logInput = true,
                  const bool full = false,
                  const typename InputDataType::elem_type eps = 1e-08,
-                 const bool mean = true);
+                 const bool reduction = true);
 
   /**
    * Computes the Poisson negative log likelihood Loss.
@@ -112,13 +116,10 @@ class PoissonNLLLoss
   //! logarithms and denominators.
   typename InputDataType::elem_type& Eps() { return eps; }
 
-  //! Get the value of mean. It's a boolean value that tells if
-  //! mean of the total loss has to be taken.
-  bool Mean() const { return mean; }
-  //! Modify the value of mean. It's a boolean value that tells if
-  //! mean of the total loss has to be taken.
-  bool& Mean() { return mean; }
-
+  //! Get the type of reduction used.
+  bool Reduction() const { return reduction; }
+  //! Modify the type of reduction used.
+  bool& Reduction() { return reduction; }
   /**
    * Serialize the layer.
    */
@@ -154,8 +155,10 @@ class PoissonNLLLoss
   //! eps is a small value required to prevent 0 in logarithms and denominators.
   typename InputDataType::elem_type eps;
 
-  //! Boolean value that tells if mean of the total loss has to be taken.
-  bool mean;
+  //! Boolean value that tells if reduction 
+  //  is 'sum' or 'mean'.
+  bool reduction;
+
 }; // class PoissonNLLLoss
 
 } // namespace ann
diff --git a/src/mlpack/methods/ann/loss_functions/poisson_nll_loss_impl.hpp b/src/mlpack/methods/ann/loss_functions/poisson_nll_loss_impl.hpp
index 05d2d79886..3152d73f56 100644
--- a/src/mlpack/methods/ann/loss_functions/poisson_nll_loss_impl.hpp
+++ b/src/mlpack/methods/ann/loss_functions/poisson_nll_loss_impl.hpp
@@ -24,11 +24,11 @@ PoissonNLLLoss<InputDataType, OutputDataType>::PoissonNLLLoss(
     const bool logInput,
     const bool full,
     const typename InputDataType::elem_type eps,
-    const bool mean):
+    const bool reduction):
     logInput(logInput),
     full(full),
     eps(eps),
-    mean(mean)
+    reduction(reduction)
 {
   Log::Assert(eps >= 0, "Epsilon (eps) must be greater than or equal to zero.");
 }
@@ -57,8 +57,12 @@ PoissonNLLLoss<InputDataType, OutputDataType>::Forward(
         + 0.5 * arma::log(2 * M_PI * target);
     loss.elem(arma::find(mask)) += approx.elem(arma::find(mask));
   }
-
-  return mean ? arma::accu(loss) / loss.n_elem : arma::accu(loss);
+  typename PredictionType::elem_type lossSum = arma::accu(loss);
+  
+  if (reduction)
+    return lossSum;
+  
+  return lossSum / loss.n_elem;
 }
 
 template<typename InputDataType, typename OutputDataType>
@@ -75,7 +79,7 @@ void PoissonNLLLoss<InputDataType, OutputDataType>::Backward(
   else
     loss = (1 - target / (prediction + eps));
 
-  if (mean)
+  if (!reduction)
     loss = loss / loss.n_elem;
 }
 
@@ -88,7 +92,7 @@ void PoissonNLLLoss<InputDataType, OutputDataType>::serialize(
   ar(CEREAL_NVP(logInput));
   ar(CEREAL_NVP(full));
   ar(CEREAL_NVP(eps));
-  ar(CEREAL_NVP(mean));
+  ar(CEREAL_NVP(reduction));
 }
 
 } // namespace ann