Use narrow instead of indexing by slices (#4014)

* Use narrow instead of indexing by slices

* Use index_select instead of select

* Move fancy indexing to python instead of c++

---------

Co-authored-by: Sam Anklesaria <[email protected]>

Author: samanklesaria

src/libtorchaudio/forced_align/cpu/compute.cpp

@@ -187,13 +187,13 @@ std::tuple<torch::Tensor, torch::Tensor> compute(
       });
   return std::make_tuple(
       paths,
-      logProbs.index(
-          {torch::indexing::Slice(),
-           torch::linspace(
-               0, T - 1, T, torch::TensorOptions().dtype(paths.dtype())),
-           paths.index({0})}));
+      logProbs
+      );
 }
 
+
+
+
 TORCH_LIBRARY_IMPL(torchaudio, CPU, m) {
   m.impl("forced_align", &compute);
 }

src/libtorchaudio/forced_align/gpu/compute.cu

@@ -306,11 +306,7 @@ std::tuple<torch::Tensor, torch::Tensor> compute(
       });
   return std::make_tuple(
       paths.to(logProbs.device()),
-      logProbs.index(
-          {torch::indexing::Slice(),
-           torch::linspace(
-               0, T - 1, T, torch::TensorOptions().dtype(paths.dtype())),
-           paths.index({0})}));
+      logProbs);
 }
 
 TORCH_LIBRARY_IMPL(torchaudio, CUDA, m) {

src/libtorchaudio/lfilter.cpp

@@ -82,15 +82,9 @@ void lfilter_core_generic_loop(
   auto coeff = a_coeff_flipped.unsqueeze(2);
   for (int64_t i_sample = 0; i_sample < n_samples_input; i_sample++) {
     auto windowed_output_signal =
-        padded_output_waveform
-            .index(
-                {torch::indexing::Slice(),
-                 torch::indexing::Slice(),
-                 torch::indexing::Slice(i_sample, i_sample + n_order)})
-            .transpose(0, 1);
+      torch::narrow(padded_output_waveform, 2, i_sample, i_sample + n_order).transpose(0, 1);
     auto o0 =
-        input_signal_windows.index(
-            {torch::indexing::Slice(), torch::indexing::Slice(), i_sample}) -
+        torch::select(input_signal_windows, 2, i_sample) -
         at::matmul(windowed_output_signal, coeff).squeeze(2).transpose(0, 1);
     padded_output_waveform.index_put_(
         {torch::indexing::Slice(),

src/torchaudio/functional/_alignment.py

@@ -70,7 +70,7 @@ def forced_align(
     assert target_lengths is not None
 
     paths, scores = torch.ops.torchaudio.forced_align(log_probs, targets, input_lengths, target_lengths, blank)
-    return paths, scores
+    return paths, scores[:, torch.arange(scores.shape[1]), paths[0]]
 
 
 @dataclass