[STABLE ABI PORT] Use stable ABI for RNNT

src/libtorchaudio/CMakeLists.txt

@@ -22,19 +22,13 @@ if(BUILD_RNNT)
   list(
     APPEND
     sources
-    rnnt/cpu/compute_alphas.cpp
-    rnnt/cpu/compute_betas.cpp
     rnnt/cpu/compute.cpp
-    rnnt/compute_alphas.cpp
-    rnnt/compute_betas.cpp
     rnnt/compute.cpp
     )
   if (USE_CUDA)
     list(
       APPEND
       sources
-      rnnt/gpu/compute_alphas.cu
-      rnnt/gpu/compute_betas.cu
       rnnt/gpu/compute.cu
       )
   endif()

src/libtorchaudio/rnnt/compute.cpp

@@ -1,34 +1,14 @@
-#include <libtorchaudio/rnnt/compute.h>
+#include <torch/csrc/stable/library.h>
+#include <torch/csrc/inductor/aoti_torch/c/shim.h>
 
-std::tuple<torch::Tensor, std::optional<torch::Tensor>> rnnt_loss(
-    torch::Tensor& logits,
-    const torch::Tensor& targets,
-    const torch::Tensor& logit_lengths,
-    const torch::Tensor& target_lengths,
-    int64_t blank,
-    double clamp,
-    bool fused_log_softmax = true) {
-  static auto op = torch::Dispatcher::singleton()
-                       .findSchemaOrThrow("torchaudio::rnnt_loss", "")
-                       .typed<decltype(rnnt_loss)>();
-  return op.call(
-      logits,
-      targets,
-      logit_lengths,
-      target_lengths,
-      blank,
-      clamp,
-      fused_log_softmax);
-}
 
-TORCH_LIBRARY_FRAGMENT(torchaudio, m) {
+STABLE_TORCH_LIBRARY_FRAGMENT(torchaudio, m) {
   m.def(
       "rnnt_loss(Tensor logits,"
       "Tensor targets,"
       "Tensor logit_lengths,"
       "Tensor target_lengths,"
       "int blank,"
       "float clamp,"
-      "bool fused_log_softmax) -> (Tensor, Tensor?)");
-  m.def("torchaudio::rnnt_loss_forward", &rnnt_loss);
+      "bool fused_log_softmax) -> (Tensor, Tensor)");
 }