2025-11-21 nightly release (32ce8c0)

Author: pytorchbot

src/libtorchaudio/forced_align/compute.cpp

@@ -1,19 +1,10 @@
-#include <libtorchaudio/forced_align/compute.h>
-#include <torch/script.h>
+#include <torch/csrc/stable/library.h>
 
-std::tuple<torch::Tensor, torch::Tensor> forced_align(
-    const torch::Tensor& logProbs,
-    const torch::Tensor& targets,
-    const torch::Tensor& inputLengths,
-    const torch::Tensor& targetLengths,
-    const int64_t blank) {
-  static auto op = torch::Dispatcher::singleton()
-                       .findSchemaOrThrow("torchaudio::forced_align", "")
-                       .typed<decltype(forced_align)>();
-  return op.call(logProbs, targets, inputLengths, targetLengths, blank);
-}
-
-TORCH_LIBRARY_FRAGMENT(torchaudio, m) {
+STABLE_TORCH_LIBRARY_FRAGMENT(torchaudio, m) {
   m.def(
-      "forced_align(Tensor log_probs, Tensor targets, Tensor input_lengths, Tensor target_lengths, int blank) -> (Tensor, Tensor)");
+      "forced_align(Tensor log_probs,"
+      "Tensor targets,"
+      "Tensor input_lengths,"
+      "Tensor target_lengths,"
+      "int blank) -> (Tensor, Tensor)");
 }

src/libtorchaudio/forced_align/compute.h

@@ -1,10 +1 @@
 #pragma once
-
-#include <torch/script.h>
-
-std::tuple<torch::Tensor, torch::Tensor> forced_align(
-    const torch::Tensor& logProbs,
-    const torch::Tensor& targets,
-    const torch::Tensor& inputLengths,
-    const torch::Tensor& targetLengths,
-    const int64_t blank);

src/libtorchaudio/forced_align/cpu/compute.cpp

@@ -38,9 +38,9 @@ void forced_align_impl(
   for (int i = 0; i < T * S; i++) {
     backPtr_a[i] = -1;
   }
-  auto logProbs_a = torchaudio::stable::accessor<scalar_t, 3>(logProbs);
-  auto targets_a = torchaudio::stable::accessor<target_t, 2>(targets);
-  auto paths_a = torchaudio::stable::accessor<target_t, 2>(paths);
+  auto logProbs_a = torchaudio::accessor<scalar_t, 3>(logProbs);
+  auto targets_a = torchaudio::accessor<target_t, 2>(targets);
+  auto paths_a = torchaudio::accessor<target_t, 2>(paths);
   auto R = 0;
   for (auto i = 1; i < L; i++) {
     if (targets_a[batchIndex][i] == targets_a[batchIndex][i - 1]) {
@@ -147,10 +147,10 @@ template <typename scalar_t>
 const auto forced_align_int_impl = forced_align_impl<scalar_t, ScalarType::Int>;
 
 std::tuple<Tensor, Tensor> compute(
-    const Tensor& logProbs,
-    const Tensor& targets,
-    const Tensor& inputLengths,
-    const Tensor& targetLengths,
+    Tensor logProbs,
+    Tensor targets,
+    Tensor inputLengths,
+    Tensor targetLengths,
     const int64_t blank) {
   STD_TORCH_CHECK(logProbs.is_cpu(), "log_probs must be a CPU tensor");
   STD_TORCH_CHECK(targets.is_cpu(), "targets must be a CPU tensor");
@@ -224,24 +224,8 @@ std::tuple<Tensor, Tensor> compute(
   return std::make_tuple(paths, logProbs);
 }
 
-void boxed_forced_align_cpu(
-    StableIValue* stack,
-    uint64_t num_args,
-    uint64_t num_outputs) {
-  STD_TORCH_CHECK(num_args == 5, "num_args must be 5");
-  STD_TORCH_CHECK(num_outputs == 2, "num_outputs must be 2");
-  std::tuple<Tensor, Tensor> res = compute(
-      /*logProbs*/ torch::stable::detail::to<Tensor>(stack[0]),
-      /*targets*/ torch::stable::detail::to<Tensor>(stack[1]),
-      /*logit_lengths*/ torch::stable::detail::to<Tensor>(stack[2]),
-      /*target_lengths*/ torch::stable::detail::to<Tensor>(stack[3]),
-      /*blank*/ float(torch::stable::detail::to<int64_t>(stack[4])));
-  stack[0] = torch::stable::detail::from(std::get<0>(res));
-  stack[1] = torch::stable::detail::from(std::get<1>(res));
-}
-
 STABLE_TORCH_LIBRARY_IMPL(torchaudio, CPU, m) {
-  m.impl("forced_align", &boxed_forced_align_cpu);
+  m.impl("forced_align", TORCH_BOX(&compute));
 }
 
 } // namespace cpu

src/libtorchaudio/forced_align/gpu/compute.cu

@@ -1,5 +1,4 @@
 #include <libtorchaudio/utils.h>
-#include <libtorchaudio/stable/TensorAccessor.h>
 #include <torch/csrc/stable/library.h>
 #include <torch/headeronly/core/Dispatch_v2.h>
 #include <torch/headeronly/core/ScalarType.h>
@@ -23,9 +22,9 @@ using torch::headeronly::ScalarType;
 
 template <typename scalar_t, typename target_t>
 __global__ void falign_cuda_step_kernel(
-    const torchaudio::stable::PackedTensorAccessor32<scalar_t, 3, torchaudio::stable::RestrictPtrTraits>
+    const torchaudio::PackedTensorAccessor32<scalar_t, 3>
         logProbs_a,
-    const torchaudio::stable::PackedTensorAccessor32<target_t, 2, torchaudio::stable::RestrictPtrTraits>
+    const torchaudio::PackedTensorAccessor32<target_t, 2>
         targets_a,
     const int T,
     const int L,
@@ -36,9 +35,9 @@ __global__ void falign_cuda_step_kernel(
     int start,
     int end,
     int backPtrBufferLen,
-    torchaudio::stable::PackedTensorAccessor32<scalar_t, 2, torchaudio::stable::RestrictPtrTraits>
+    torchaudio::PackedTensorAccessor32<scalar_t, 2>
         alphas_a,
-    torchaudio::stable::PackedTensorAccessor32<int8_t, 2, torchaudio::stable::RestrictPtrTraits>
+    torchaudio::PackedTensorAccessor32<int8_t, 2>
         backPtrBuffer_a) {
   scalar_t kNegInfinity = -std::numeric_limits<scalar_t>::infinity();
   const int batchIndex =
@@ -125,7 +124,7 @@ void forced_align_impl(
   const scalar_t kNegInfinity = -std::numeric_limits<scalar_t>::infinity();
   using target_t = typename std::
       conditional<target_scalar_type == ScalarType::Int, int, int64_t>::type;
-  auto paths_a = torchaudio::stable::accessor<target_t, 2>(paths);
+  auto paths_a = torchaudio::accessor<target_t, 2>(paths);
   const int batchIndex =
       0; // TODO: support batch version and use the real batch index
   const int T = logProbs.size(1); // num frames
@@ -150,8 +149,8 @@ void forced_align_impl(
   torch::stable::fill_(alphas, kNegInfinity);
 
   // CPU accessors
-  auto targetsCpu_a = torchaudio::stable::accessor<target_t, 2>(targetsCpu);
-  auto backPtrCpu_a = torchaudio::stable::accessor<int8_t, 2>(backPtrCpu);
+  auto targetsCpu_a = torchaudio::accessor<target_t, 2>(targetsCpu);
+  auto backPtrCpu_a = torchaudio::accessor<int8_t, 2>(backPtrCpu);
   // count the number of repeats in label
   int R = 0;
   for (int i = 1; i < L; ++i) {
@@ -192,8 +191,8 @@ void forced_align_impl(
     }
     falign_cuda_step_kernel<scalar_t, target_t>
         <<<1, kNumThreads, 0, defaultStream>>>(
-            torchaudio::stable::packed_accessor32<scalar_t, 3, torchaudio::stable::RestrictPtrTraits>(logProbs),
-            torchaudio::stable::packed_accessor32<target_t, 2, torchaudio::stable::RestrictPtrTraits>(targets),
+            torchaudio::packed_accessor32<scalar_t, 3>(logProbs),
+            torchaudio::packed_accessor32<target_t, 2>(targets),
             T,
             L,
             N,
@@ -203,8 +202,8 @@ void forced_align_impl(
             start,
             end,
             backPtrBufferLen,
-            torchaudio::stable::packed_accessor32<scalar_t, 2, torchaudio::stable::RestrictPtrTraits>(alphas),
-            torchaudio::stable::packed_accessor32<int8_t, 2, torchaudio::stable::RestrictPtrTraits>(backPtrBuffer));
+            torchaudio::packed_accessor32<scalar_t, 2>(alphas),
+            torchaudio::packed_accessor32<int8_t, 2>(backPtrBuffer));
     C10_CUDA_KERNEL_LAUNCH_CHECK();
     ++backPtrBufferLen;
     if (backPtrBufferLen == kBackPtrBufferSize || t == T - 1) {
@@ -228,9 +227,8 @@ void forced_align_impl(
     }
   }
   cpuDataTranferStream.synchronize();
-
   auto alphasCpu = torchaudio::stable::cpu(alphas);
-  auto alphasCpu_a = torchaudio::stable::accessor<scalar_t, 2>(alphasCpu);
+  auto alphasCpu_a = torchaudio::accessor<scalar_t, 2>(alphasCpu);
   int curIdxOffset = ((T - 1) % 2);
   int ltrIdx =
       alphasCpu_a[curIdxOffset][S - 1] > alphasCpu_a[curIdxOffset][S - 2]
@@ -244,18 +242,11 @@ void forced_align_impl(
   }
 }
 
-template <typename scalar_t>
-const auto forced_align_long_impl =
-    forced_align_impl<scalar_t, ScalarType::Long>;
-
-template <typename scalar_t>
-const auto forced_align_int_impl = forced_align_impl<scalar_t, ScalarType::Int>;
-
 std::tuple<Tensor, Tensor> compute(
-    const Tensor& logProbs,
-    const Tensor& targets,
-    const Tensor& inputLengths,
-    const Tensor& targetLengths,
+    Tensor logProbs,
+    Tensor targets,
+    Tensor inputLengths,
+    Tensor targetLengths,
     const int64_t blank) {
 
   STD_TORCH_CHECK(logProbs.is_cuda(), "log_probs must be a CUDA tensor");
@@ -307,31 +298,17 @@ std::tuple<Tensor, Tensor> compute(
 
   THO_DISPATCH_V2(logProbs.scalar_type(), "forced_align_impl", AT_WRAP([&] {
         if (targets.scalar_type() == ScalarType::Long) {
-          forced_align_long_impl<scalar_t>(logProbs, targets, blank, paths);
+          (forced_align_impl<scalar_t, ScalarType::Long>(logProbs, targets, blank, paths));
         } else {
-          forced_align_int_impl<scalar_t>(logProbs, targets, blank, paths);
-        }
+          (forced_align_impl<scalar_t, ScalarType::Int>(logProbs, targets, blank, paths));
+          }
       }), AT_EXPAND(AT_FLOATING_TYPES), ScalarType::Half);
-
   Tensor pathsCuda = torchaudio::stable::cuda(paths, logProbs.get_device_index());
   return std::make_tuple(pathsCuda, logProbs);
 }
 
-void boxed_forced_align_gpu(StableIValue* stack, uint64_t num_args, uint64_t num_outputs) {
-  STD_TORCH_CHECK(num_args == 5, "num_args must be 5");
-  STD_TORCH_CHECK(num_outputs == 2, "num_outputs must be 2");
-  std::tuple<Tensor, Tensor> res = compute(
-      /*logProbs*/torch::stable::detail::to<Tensor>(stack[0]),
-      /*targets*/torch::stable::detail::to<Tensor>(stack[1]),
-      /*logit_lengths*/torch::stable::detail::to<Tensor>(stack[2]),
-      /*target_lengths*/torch::stable::detail::to<Tensor>(stack[3]),
-      /*blank*/float(torch::stable::detail::to<int64_t>(stack[4])));
-  stack[0] = torch::stable::detail::from(std::get<0>(res));
-  stack[1] = torch::stable::detail::from(std::get<1>(res));
-}
-
 STABLE_TORCH_LIBRARY_IMPL(torchaudio, CUDA, m) {
-  m.impl("forced_align", &boxed_forced_align_gpu);
+  m.impl("forced_align", TORCH_BOX(&compute));
 }
 
 } // namespace gpu

src/libtorchaudio/overdrive.cpp

@@ -1,52 +1,71 @@
-#include <torch/script.h>
-#include <torch/torch.h>
+#include <libtorchaudio/utils.h>
+#include <torch/csrc/stable/library.h>
+#include <torch/csrc/stable/ops.h>
+#include <torch/csrc/stable/tensor.h>
+#include <torch/headeronly/core/Dispatch_v2.h>
+#include <torch/headeronly/core/TensorAccessor.h>
 
 namespace {
+using torch::stable::Tensor;
 
 template <typename scalar_t>
 void overdrive_cpu_kernel(
-    at::TensorAccessor<scalar_t, 2> waveform_accessor,
-    at::TensorAccessor<scalar_t, 2> temp_accessor,
-    at::TensorAccessor<scalar_t, 1> last_in_accessor,
-    at::TensorAccessor<scalar_t, 1> last_out_accessor,
-    at::TensorAccessor<scalar_t, 2> output_waveform_accessor) {
+    torchaudio::TensorAccessor<scalar_t, 2> waveform_accessor,
+    torchaudio::TensorAccessor<scalar_t, 2> temp_accessor,
+    torchaudio::TensorAccessor<scalar_t, 1> last_in_accessor,
+    torchaudio::TensorAccessor<scalar_t, 1> last_out_accessor,
+    torchaudio::TensorAccessor<scalar_t, 2> output_waveform_accessor) {
   int64_t n_frames = waveform_accessor.size(1);
   int64_t n_channels = waveform_accessor.size(0);
 
-  at::parallel_for(0, n_channels, 1, [&](int64_t begin, int64_t end) {
-    for (int64_t i_channel = begin; i_channel < end; ++i_channel) {
-      for (int64_t i_frame = 0; i_frame < n_frames; ++i_frame) {
-        last_out_accessor[i_channel] = temp_accessor[i_channel][i_frame] -
-            last_in_accessor[i_channel] + 0.995 * last_out_accessor[i_channel];
-        last_in_accessor[i_channel] = temp_accessor[i_channel][i_frame];
-        output_waveform_accessor[i_channel][i_frame] =
-            waveform_accessor[i_channel][i_frame] * 0.5 +
-            last_out_accessor[i_channel] * 0.75;
-      }
-    }
-  });
+  torch::stable::parallel_for(
+      0, n_channels, 1, [&](int64_t begin, int64_t end) {
+        for (int64_t i_channel = begin; i_channel < end; ++i_channel) {
+          for (int64_t i_frame = 0; i_frame < n_frames; ++i_frame) {
+            last_out_accessor[i_channel] = temp_accessor[i_channel][i_frame] -
+                last_in_accessor[i_channel] +
+                0.995 * last_out_accessor[i_channel];
+            last_in_accessor[i_channel] = temp_accessor[i_channel][i_frame];
+            output_waveform_accessor[i_channel][i_frame] =
+                waveform_accessor[i_channel][i_frame] * 0.5 +
+                last_out_accessor[i_channel] * 0.75;
+          }
+        }
+      });
 }
 
-void overdrive_core_loop_cpu(
-    at::Tensor& waveform,
-    at::Tensor& temp,
-    at::Tensor& last_in,
-    at::Tensor& last_out,
-    at::Tensor& output_waveform) {
-  AT_DISPATCH_FLOATING_TYPES(waveform.scalar_type(), "overdrive_cpu", ([&] {
-                               overdrive_cpu_kernel<scalar_t>(
-                                   waveform.accessor<scalar_t, 2>(),
-                                   temp.accessor<scalar_t, 2>(),
-                                   last_in.accessor<scalar_t, 1>(),
-                                   last_out.accessor<scalar_t, 1>(),
-                                   output_waveform.accessor<scalar_t, 2>());
-                             }));
+std::tuple<Tensor, Tensor, Tensor> overdrive_core_loop_cpu(
+    Tensor waveform,
+    Tensor temp,
+    Tensor last_in,
+    Tensor last_out,
+    Tensor output_waveform) {
+  THO_DISPATCH_V2(
+      waveform.scalar_type(),
+      "overdrive_cpu",
+      AT_WRAP([&] {
+        overdrive_cpu_kernel<scalar_t>(
+            torchaudio::accessor<scalar_t, 2>(waveform),
+            torchaudio::accessor<scalar_t, 2>(temp),
+            torchaudio::accessor<scalar_t, 1>(last_in),
+            torchaudio::accessor<scalar_t, 1>(last_out),
+            torchaudio::accessor<scalar_t, 2>(output_waveform));
+      }),
+      AT_FLOATING_TYPES);
+  return std::make_tuple(last_in, last_out, output_waveform);
 }
 
 } // namespace
 
-// Note: We want to avoid using "catch-all" kernel.
-// The following registration should be replaced with CPU specific registration.
-TORCH_LIBRARY_FRAGMENT(torchaudio, m) {
-  m.def("torchaudio::_overdrive_core_loop", &overdrive_core_loop_cpu);
+STABLE_TORCH_LIBRARY_FRAGMENT(torchaudio, m) {
+  m.def(
+      "_overdrive_core_loop(Tensor waveform,"
+      "Tensor temp,"
+      "Tensor(a!) last_in,"
+      "Tensor(b!) last_out,"
+      "Tensor(c!) output_waveform) -> (Tensor(a!), Tensor(b!), Tensor(c!))");
+}
+
+STABLE_TORCH_LIBRARY_IMPL(torchaudio, CPU, m) {
+  m.impl("_overdrive_core_loop", TORCH_BOX(&overdrive_core_loop_cpu));
 }