Shuffle sequence lengths (#3926)

COPYRIGHT.txt

@@ -162,6 +162,7 @@ Copyright:
   Copyright 2024, Felix Patschkowski <felix.patschkowski@gmail.com>
   Copyright 2025, Benjamin A. Beasley <code@musicinmybrain.net>
   Copyright 2025, Maksym Prots <imaxprots@gmail.com>
+  Copyright 2025, Zachary Ng <zachn716@gmail.com>
 
 License: BSD-3-clause
   All rights reserved.

HISTORY.md

@@ -3,6 +3,8 @@
 ## mlpack ?.?.?
 
 _????-??-??_
+ * Shuffle sequence lengths (#3926)
+
  * Add ability to compile OpenBLAS for windows (#3922)
 
  * Drop pytest-runner and "setup.py test" support (#3921).

src/mlpack/core/math/shuffle_data.hpp

@@ -146,6 +146,74 @@ void ShuffleData(const MatType& inputPoints,
   }
 }
 
+/**
+ * Shuffle a cube-shaped dataset and associated labels (or responses) which are
+ * also cube-shaped.  Also shuffle its weights.  It is expected that inputPoints, 
+ * inputLabels, and inputWeights have the same number of columns.
+ *
+ * Shuffled data will be output into outputPoints, outputLabels, and outputWeights.
+ */
+template<typename MatType, typename LabelsType, typename WeightsType>
+void ShuffleData(const MatType& inputPoints,
+                 const LabelsType& inputLabels,
+                 const WeightsType& inputWeights,
+                 MatType& outputPoints,
+                 LabelsType& outputLabels,
+                 WeightsType& outputWeights,
+                 const std::enable_if_t<!arma::is_SpMat<MatType>::value>* = 0,
+                 const std::enable_if_t<arma::is_Cube<MatType>::value>* = 0,
+                 const std::enable_if_t<arma::is_Cube<LabelsType>::value>* = 0)
+{
+  // Generate ordering.
+  arma::uvec ordering = arma::shuffle(arma::linspace<arma::uvec>(0,
+      inputPoints.n_cols - 1, inputPoints.n_cols));
+
+  // Properly handle the case where the input and output data are the same
+  // object.
+  MatType* outputPointsPtr = &outputPoints;
+  LabelsType* outputLabelsPtr = &outputLabels;
+  WeightsType* outputWeightsPtr = &outputWeights;
+  if (&inputPoints == &outputPoints)
+    outputPointsPtr = new MatType();
+  if (&inputLabels == &outputLabels)
+    outputLabelsPtr = new LabelsType();
+  if (&inputWeights == &outputWeights)
+    outputWeightsPtr = new WeightsType();
+
+  outputPointsPtr->set_size(inputPoints.n_rows, inputPoints.n_cols,
+      inputPoints.n_slices);
+  outputLabelsPtr->set_size(inputLabels.n_rows, inputLabels.n_cols,
+      inputLabels.n_slices);
+  outputWeightsPtr->set_size(inputWeights.n_cols);
+  for (size_t i = 0; i < ordering.n_elem; ++i)
+  {
+    outputPointsPtr->tube(0, ordering[i], outputPointsPtr->n_rows - 1,
+        ordering[i]) = inputPoints.tube(0, i, inputPoints.n_rows - 1, i);
+    outputLabelsPtr->tube(0, ordering[i], outputLabelsPtr->n_rows - 1,
+        ordering[i]) = inputLabels.tube(0, i, inputLabels.n_rows - 1, i);
+    outputWeightsPtr->at(ordering[i]) = inputWeights[i];
+  }
+
+  // Clean up memory if needed.
+  if (&inputPoints == &outputPoints)
+  {
+    outputPoints = std::move(*outputPointsPtr);
+    delete outputPointsPtr;
+  }
+
+  if (&inputLabels == &outputLabels)
+  {
+    outputLabels = std::move(*outputLabelsPtr);
+    delete outputLabelsPtr;
+  }
+
+  if (&inputWeights == &outputWeights)
+  {
+    outputWeights = std::move(*outputWeightsPtr);
+    delete outputWeightsPtr;
+  }
+}
+
 /**
  * Shuffle a dataset and associated labels (or responses) and weights.  It is
  * expected that inputPoints and inputLabels and inputWeights have the same

src/mlpack/methods/ann/rnn_impl.hpp

@@ -633,7 +633,15 @@ void RNN<
     MatType
 >::Shuffle()
 {
-  ShuffleData(predictors, responses, predictors, responses);
+  if (sequenceLengths.n_elem > 0)
+  {
+    ShuffleData(predictors, responses, sequenceLengths,
+        predictors, responses, sequenceLengths);
+  }
+  else
+  {
+    ShuffleData(predictors, responses, predictors, responses);
+  }
 }
 
 template<

src/mlpack/tests/math_test.cpp

@@ -696,6 +696,75 @@ TEST_CASE("CubeShuffleTest", "[MathTest]")
     REQUIRE(counts[i] == data.n_slices);
 }
 
+/**
+ * Make sure shuffling cubes with ragged sequence lengths works.
+ */
+TEST_CASE("RaggedCubeShuffleTest", "[MathTest]")
+{
+  arma::cube data(3, 5, 5);
+  arma::cube labels(1, 5, 5);
+  arma::Row<size_t> lengths(5);
+
+  data.fill(-1);
+  labels.fill(-1);
+
+  for (size_t c = 0; c < lengths.n_elem; ++c)
+  {
+    lengths[c] = c;
+    for (size_t s = 0; s < lengths[c]; ++s)
+    {
+      data(0, c, s) = s;
+      data(1, c, s) = c;
+      labels(0, c, s) = c + s;
+    }
+  }
+
+  arma::cube outputData, outputLabels;
+  arma::Row<size_t> outputLengths;
+
+  ShuffleData(data, labels, lengths, outputData, outputLabels, outputLengths);
+
+  REQUIRE(outputData.n_rows == data.n_rows);
+  REQUIRE(outputData.n_cols == data.n_cols);
+  REQUIRE(outputData.n_slices == data.n_slices);
+  REQUIRE(outputLabels.n_rows == labels.n_rows);
+  REQUIRE(outputLabels.n_cols == labels.n_cols);
+  REQUIRE(outputLabels.n_slices == labels.n_slices);
+  REQUIRE(lengths.n_elem == outputLengths.n_elem);
+
+  // Make sure each column has the right number of slices
+  arma::Row<size_t> sliceCount(5);
+  for (size_t i = 0; i < outputLabels.n_cols; ++i)
+  {
+    for (size_t j = 0; j < outputLabels.n_slices; j++)
+    {
+      if (outputLabels(0, i, j) < 0) {
+        sliceCount[i] = j;
+        break;
+      }
+    }
+  }
+
+  for (size_t i = 0; i < 5; ++i)
+    REQUIRE(sliceCount[i] == outputLengths[i]);
+
+  // Make sure we only have each point once.
+  arma::Row<size_t> counts(5);
+  for (size_t c = 0; c < 5; ++c)
+  {
+    for (size_t s = 0; s < outputLengths[c]; ++s)
+    {
+      REQUIRE(outputData(0, c, s) + outputData(1, c, s)
+          == outputLabels(0, c, s));
+      REQUIRE(outputData(2, c, s) == Approx(-1.0).margin(1e-5));
+      counts[outputLengths[c]]++;
+    }
+  }
+
+  for (size_t i = 0; i < 5; ++i)
+    REQUIRE(counts[i] == i);
+}
+
 /**
  * Make sure shuffling data with weights works.
  */