From b4d2adbc677902a23210e99b29f773f89cfb07e4 Mon Sep 17 00:00:00 2001
From: Ryan Curtin <ryan@ratml.org>
Date: Sat, 11 Jun 2022 10:53:16 -0400
Subject: [PATCH] Some minor bugfixes (but not totally fixed yet).

---
 src/mlpack/methods/ann/layer/concat.hpp      | 22 ++++++++++++++++----
 src/mlpack/methods/ann/layer/concat_impl.hpp | 15 ++++++-------
 2 files changed, 26 insertions(+), 11 deletions(-)

diff --git a/src/mlpack/methods/ann/layer/concat.hpp b/src/mlpack/methods/ann/layer/concat.hpp
index 7d8aac0f3d..da17984c2d 100644
--- a/src/mlpack/methods/ann/layer/concat.hpp
+++ b/src/mlpack/methods/ann/layer/concat.hpp
@@ -159,17 +159,18 @@ class ConcatType : public MultiLayer<MatType>
     }
 
     // Now, we concatenate the output along a specific axis.
-    this->outputDimensions = std::vector<size_t>(this->inputDimensions.size(),
+    this->outputDimensions = std::vector<size_t>(
+        (this->network.size() == 0) ?
+            this->inputDimensions.size() :
+            this->network[0]->OutputDimensions().size(),
         0);
-    for (size_t i = 0; i < this->inputDimensions.size(); ++i)
+    for (size_t i = 0; i < this->outputDimensions.size(); ++i)
     {
       if (i == axis)
       {
         // Accumulate output size along this axis for each layer output.
         for (size_t n = 0; n < this->network.size(); ++n)
-        {
           this->outputDimensions[i] += this->network[n]->OutputDimensions()[i];
-        }
       }
       else
       {
@@ -194,6 +195,19 @@ class ConcatType : public MultiLayer<MatType>
         this->outputDimensions[i] = axisDim;
       }
     }
+
+    // Recompute total input and output sizes.  Note that we pass the input to
+    // each layer held in the network, so the "total" input size (which is used
+    // by the backwards pass to compute how much memory to use for holding
+    // deltas) should be the number of layers multiplied by the input size for
+    // each layer.
+    this->totalInputSize = 1;
+    this->totalOutputSize = 1;
+    for (size_t i = 0; i < this->inputDimensions.size(); ++i)
+      this->totalInputSize *= this->inputDimensions[i];
+    this->totalInputSize *= this->network.size();
+    for (size_t i = 0; i < this->outputDimensions.size(); ++i)
+      this->totalOutputSize *= this->outputDimensions[i];
   }
 
   /**
diff --git a/src/mlpack/methods/ann/layer/concat_impl.hpp b/src/mlpack/methods/ann/layer/concat_impl.hpp
index 29c345492b..c4710fd79e 100644
--- a/src/mlpack/methods/ann/layer/concat_impl.hpp
+++ b/src/mlpack/methods/ann/layer/concat_impl.hpp
@@ -112,18 +112,19 @@ void ConcatType<MatType>::Forward(const MatType& input, MatType& output)
   // this->outputDimensions.size(); that is the batch size (represented as the
   // number of columns in `input`).
 
-  size_t slices = (axis == 0) ? input.n_cols :
-      std::accumulate(this->outputDimensions.begin(),
-          this->outputDimensions.begin() + axis, 0) + input.n_cols;
-  size_t rows = (axis == this->outputDimensions.size() - 1) ? 1 :
-      std::accumulate(this->outputDimensions.begin() + axis + 1,
-          this->outputDimensions.end(), 0);
+  size_t rows = 1;
+  for (size_t i = 0; i < axis; ++i)
+    rows *= this->outputDimensions[i];
+
+  size_t slices = 1;
+  for (size_t i = axis + 1; i < this->outputDimensions.size(); ++i)
+    slices *= this->outputDimensions[i];
 
   std::vector<arma::Cube<typename MatType::elem_type>> layerOutputAliases(
       this->layerOutputs.size());
   for (size_t i = 0; i < this->layerOutputs.size(); ++i)
   {
-    MakeAlias(layerOutputAliases.back(),
+    MakeAlias(layerOutputAliases[i],
               (typename MatType::elem_type*) this->layerOutputs[i].memptr(),
               rows,
               this->network[i]->OutputDimensions()[axis],