Aggregate TORCH_LIBRARY calls

Author: samanklesaria

src/libtorchaudio/CMakeLists.txt

@@ -25,8 +25,6 @@ if(BUILD_RNNT)
     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)

src/libtorchaudio/rnnt/compute.cpp

@@ -9,4 +9,18 @@ STABLE_TORCH_LIBRARY(torchaudio, m) {
       "int blank,"
       "float clamp,"
       "bool fused_log_softmax) -> (Tensor, Tensor?)");
+  m.def(
+      "rnnt_loss_betas(Tensor logits,"
+      "Tensor targets,"
+      "Tensor logit_lengths,"
+      "Tensor target_lengths,"
+      "int blank,"
+      "float clamp) -> Tensor");
+  m.def(
+      "rnnt_loss_alphas(Tensor logits,"
+      "Tensor targets,"
+      "Tensor logit_lengths,"
+      "Tensor target_lengths,"
+      "int blank,"
+      "float clamp) -> Tensor");
 }

src/libtorchaudio/rnnt/compute_alphas.cpp

@@ -210,7 +210,7 @@ void boxed_compute(StableIValue* stack, uint64_t num_args, uint64_t num_outputs)
 }
 
 STABLE_TORCH_LIBRARY_IMPL(torchaudio, CUDA, m) {
-  m.impl("rnnt_loss", &boxed_compute);
+  m.impl("torchaudio::rnnt_loss", &boxed_compute);
 }
 
 } // namespace gpu

src/libtorchaudio/rnnt/compute_betas.cpp

@@ -1763,7 +1763,7 @@ def _fix_waveform_shape(
 class RnntLoss(torch.autograd.Function):
     @staticmethod
     def forward(ctx, *args):
-        output, saved = torch.ops.torchaudio.rnnt_loss_forward(*args)
+        output, saved = torch.ops.torchaudio.rnnt_loss(*args)
         ctx.save_for_backward(saved)
         return output