Merge branch 'forced_align_backptr' into forced_align_accessors

Author: samanklesaria

src/libtorchaudio/forced_align/cpu/compute.cpp

@@ -32,10 +32,9 @@ void forced_align_impl(
   const auto L = targets.size(1);
   const auto S = 2 * L + 1;
 
-  auto alphas_a = new scalar_t[S][2]; // scalar_t is just logProbs.dtype()
-  for (int i = 0; i < S; i++) {
-    alphas_a[i][0] = kNegInfinity;
-    alphas_a[i][1] = kNegInfinity;
+  auto alphas_a = new scalar_t[2 * S]; // scalar_t is just logProbs.dtype()
+  for (int i = 0; i < 2 * S; i++) {
+    alphas_a[i] = kNegInfinity;
   }
 
   auto backPtr_a = new int8_t[T * S];
@@ -64,7 +63,8 @@ void forced_align_impl(
   auto end = (S == 1) ? 1 : 2;
   for (auto i = start; i < end; i++) {
     auto labelIdx = (i % 2 == 0) ? blank : targets_a.index(batchIndex, i / 2);
-    alphas_a[i][0] = logProbs_a.index(batchIndex,0,labelIdx);
+    alphas_a[i] = logProbs_a.index(batchIndex,0,labelIdx);
+
   }
   for (auto t = 1; t < T; t++) {
     if (T - t <= L + R) {
@@ -87,18 +87,18 @@ void forced_align_impl(
     auto curIdxOffset = t % 2;
     auto prevIdxOffset = (t - 1) % 2;
     for (auto j = 0; j < S; ++j) {
-      alphas_a[j][curIdxOffset] = -std::numeric_limits<scalar_t>::infinity();
+      alphas_a[curIdxOffset * S + j] = -std::numeric_limits<scalar_t>::infinity(); // alphas_a[curIdxOffset][j]
     }
     if (start == 0) {
-      alphas_a[0][curIdxOffset] =
-          alphas_a[0][prevIdxOffset] + logProbs_a.index(batchIndex, t, blank);
+      alphas_a[curIdxOffset * S] =
+          alphas_a[prevIdxOffset * S] + logProbs_a.index(batchIndex, t, blank);
       backPtr_a[S * t] = 0;  // backPtr_a[t][0] = 0
       startloop += 1;
     }
 
     for (auto i = startloop; i < end; i++) {
-      auto x0 = alphas_a[i][prevIdxOffset];
-      auto x1 = alphas_a[i - 1][prevIdxOffset];
+      auto x0 = alphas_a[prevIdxOffset * S + i]; // alphas_a[prevIdxOffset][i];
+      auto x1 = alphas_a[prevIdxOffset * S + i - 1]; // alphas_a[prevIdxOffset][i - 1];
       auto x2 = -std::numeric_limits<scalar_t>::infinity();
 
       auto labelIdx = (i % 2 == 0) ? blank : targets_a.index(batchIndex, i / 2);
@@ -109,7 +109,7 @@ void forced_align_impl(
       // (i != 1) just ensures we don't access targets[i - 2] if its i < 2
       if (i % 2 != 0 && i != 1 &&
           targets_a.index(batchIndex, i / 2) != targets_a.index(batchIndex, i / 2 - 1)) {
-        x2 = alphas_a[i - 2][prevIdxOffset];
+        x2 = alphas_a[prevIdxOffset * S + i - 2]; // alphas_a[prevIdxOffset][i - 2];
       }
       scalar_t result = 0.0;
       if (x2 > x1 && x2 > x0) {
@@ -122,11 +122,13 @@ void forced_align_impl(
         result = x0;
         backPtr_a[t * S + i] = 0; // backPtr_a[t][i] = 0
       }
-      alphas_a[i][curIdxOffset] = result + logProbs_a.index(batchIndex, t, labelIdx);
+
+      alphas_a[curIdxOffset * S + i] = result + logProbs_a.index(batchIndex, t, labelIdx); // alphas_a[curIdxOffset][i]
     }
   }
   auto idx1 = (T - 1) % 2;
-  auto ltrIdx = alphas_a[S - 1][idx1] > alphas_a[S - 2][idx1] ? S - 1 : S - 2;
+  auto ltrIdx = alphas_a[S * idx1 + S - 1] >
+    alphas_a[S * idx1 + S - 2] ? S - 1 : S - 2; // alphas_a[idx1][S - 1], alphas_a[idx1][S - 2]
   delete[] alphas_a;
   // path stores the token index for each time step after force alignment.
   for (auto t = T - 1; t > -1; t--) {