Use stable ABI for RNNT

.github/scripts/unittest-linux/run_test.sh

@@ -30,5 +30,5 @@ fi
 
 (
     cd test
-    pytest torchaudio_unittest -k "not backend and not /io/ and not prototype and not sox and not ffmpeg and not fairseq and not hdemucs"
+    pytest torchaudio_unittest -k "not backend and not /io/ and not prototype and not sox and not ffmpeg and not fairseq and not hdemucs and not (torchscript and rnnt)"
 )

src/libtorchaudio/CMakeLists.txt

@@ -28,7 +28,6 @@ if(BUILD_RNNT)
     rnnt/compute_alphas.cpp
     rnnt/compute_betas.cpp
     rnnt/compute.cpp
-    rnnt/autograd.cpp
     )
   if (USE_CUDA)
     list(

src/libtorchaudio/rnnt/compute.cpp

@@ -1,33 +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?)");
+      "bool fused_log_softmax) -> (Tensor, Tensor)");
 }

src/libtorchaudio/rnnt/compute_alphas.cpp

@@ -1,6 +1,6 @@
-#include <torch/script.h>
+#include <torch/csrc/stable/library.h>
 
-TORCH_LIBRARY_FRAGMENT(torchaudio, m) {
+STABLE_TORCH_LIBRARY_FRAGMENT(torchaudio, m) {
   m.def(
       "rnnt_loss_alphas(Tensor logits,"
       "Tensor targets,"

src/libtorchaudio/rnnt/compute_betas.cpp

@@ -1,6 +1,6 @@
-#include <torch/script.h>
+#include <torch/csrc/stable/library.h>
 
-TORCH_LIBRARY_FRAGMENT(torchaudio, m) {
+STABLE_TORCH_LIBRARY_FRAGMENT(torchaudio, m) {
   m.def(
       "rnnt_loss_betas(Tensor logits,"
       "Tensor targets,"