Skip to content

Commit

Permalink
Add aten::_ctc_loss and its variants (#925)
Browse files Browse the repository at this point in the history
- [x] _ctc_loss
- [x] _ctc_loss.Tensor
- [x]  _ctc_loss_backward
- [x] _ctc_loss_backward.Tensor
  • Loading branch information
xytintel authored Oct 13, 2024
1 parent b539b8d commit 9094e43
Show file tree
Hide file tree
Showing 6 changed files with 1,363 additions and 2 deletions.
43 changes: 43 additions & 0 deletions src/ATen/native/xpu/LossCTC.cpp
Original file line number Diff line number Diff line change
@@ -0,0 +1,43 @@
#include <ATen/core/Reduction.h>
#include <ATen/core/Tensor.h>
#include <ATen/native/xpu/sycl/LossCTCKernels.h>
#include <comm/RegisterUtils.h>

namespace at {
namespace native {

std::tuple<Tensor, Tensor> ctc_loss_xpu(
const Tensor& log_probs,
const Tensor& targets,
IntArrayRef input_lengths,
IntArrayRef target_lengths,
int64_t blank,
bool zero_infinity) {
return native::xpu::ctc_loss_kernel(
log_probs, targets, input_lengths, target_lengths, blank, zero_infinity);
}

Tensor ctc_loss_backward_xpu(
const Tensor& grad,
const Tensor& log_probs,
const Tensor& targets,
IntArrayRef input_lengths,
IntArrayRef target_lengths,
const Tensor& neg_log_likelihood,
const Tensor& log_alpha,
int64_t blank,
bool zero_infinity) {
return native::xpu::ctc_loss_backward_kernel(
grad,
log_probs,
targets,
input_lengths,
target_lengths,
neg_log_likelihood,
log_alpha,
blank,
zero_infinity);
}

} // namespace native
} // namespace at
2 changes: 0 additions & 2 deletions src/ATen/native/xpu/XPUFallback.template
Original file line number Diff line number Diff line change
Expand Up @@ -165,8 +165,6 @@ TORCH_LIBRARY_IMPL(aten, XPU, m) {
"cholesky",
"cholesky_inverse",
"_cholesky_solve_helper",
"_ctc_loss",
"_ctc_loss_backward",
"_cummax_helper",
"_cummin_helper",
"dot",
Expand Down
Loading

0 comments on commit 9094e43

Please sign in to comment.