From 3f0e2feccde76edd3955727b10cd2e30af93af83 Mon Sep 17 00:00:00 2001
From: Konstantin Sidorov <kost2906@yandex.ru>
Date: Fri, 21 Jul 2017 19:29:46 +0300
Subject: [PATCH] Fixes from @rcurtin's review

---
 src/mlpack/methods/ann/augmented/tasks/add.hpp      | 5 +++--
 src/mlpack/methods/ann/augmented/tasks/add_impl.hpp | 3 +++
 src/mlpack/methods/ann/augmented/tasks/sort.hpp     | 6 ++++--
 3 files changed, 10 insertions(+), 4 deletions(-)

diff --git a/src/mlpack/methods/ann/augmented/tasks/add.hpp b/src/mlpack/methods/ann/augmented/tasks/add.hpp
index 668a034e89..9f563c0b4d 100644
--- a/src/mlpack/methods/ann/augmented/tasks/add.hpp
+++ b/src/mlpack/methods/ann/augmented/tasks/add.hpp
@@ -26,8 +26,9 @@ namespace tasks /* Task utilities for augmented */ {
  * The parameters are:
  * - macimum binary length;
  * 
- * Every element of sequence is encoded as 3-dimensional vector
- * (one-hot vector for elements {0, 1, +}).
+ * Every element of sequence is encoded as 1-dimensional vector
+ * (possible vector elements are {0, 1, 0.5} -
+ * the latter corresponds to '+' sign').
  * Generated datasets are compliant with mlpack format -
  * every dataset element is shaped as a vector of
  * length 3 * (sequence length),
diff --git a/src/mlpack/methods/ann/augmented/tasks/add_impl.hpp b/src/mlpack/methods/ann/augmented/tasks/add_impl.hpp
index c1b3e70807..179650a913 100644
--- a/src/mlpack/methods/ann/augmented/tasks/add_impl.hpp
+++ b/src/mlpack/methods/ann/augmented/tasks/add_impl.hpp
@@ -54,6 +54,9 @@ const void AddTask::Generate(arma::field<arma::mat>& input,
     input(i) = arma::randi<arma::mat>(sizeA + sizeB + 1,
                                       1,
                                       arma::distr_param(0, 1));
+    // Adding leading 1 to make the distribution over numbers uniform.
+    input(i).at(sizeA - 1, 0) = 1;
+    input(i).at(sizeA + sizeB, 0) = 1;
     // Insert special value for '+' delimiter.
     labels(i) = arma::zeros(sizeA + sizeB + 1, 1);
     input(i).at(sizeA, 0) = 0.5;
diff --git a/src/mlpack/methods/ann/augmented/tasks/sort.hpp b/src/mlpack/methods/ann/augmented/tasks/sort.hpp
index 877fbb21d5..be88c06035 100644
--- a/src/mlpack/methods/ann/augmented/tasks/sort.hpp
+++ b/src/mlpack/methods/ann/augmented/tasks/sort.hpp
@@ -32,8 +32,10 @@ namespace tasks /* Task utilities for augmented */ {
  * length (binary length) * (sequence length).
  * 
  * Example of generated dataset (sequence length = 3, binary length = 2):
- * - Input sequences: [ [1,1,0,0,0,1], [0,0,0,0,1,1] ]
- * - Output sequences: [ [0,0,0,1,1,1], [0,0,0,0,1,1] ]
+ * - Input sequences: [1,1,0,0,0,1]
+ * (three numbers in the sequence are 11, 00, and 01)
+ * - Output sequences: [0,0,0,1,1,1]
+ * (00, 01, 11 - reordering of the numbers above in the ascending order)
  * 
  */
 class SortTask