Merge branch 'minjean/welford_layernorm' of https://github.com/intel/…

…torch-xpu-ops into minjean/welford_layernorm
intel · Feb 19, 2025 · 5fe435b · 5fe435b
2 parents 4358e06 + e3a1135
commit 5fe435b
Show file tree

Hide file tree

Showing 10 changed files with 664 additions and 242 deletions.
diff --git a/src/ATen/native/sparse/SparseCsrTensor.cpp b/src/ATen/native/sparse/SparseCsrTensor.cpp
diff --git a/src/ATen/native/sparse/xpu/SparseCsrTensorMath.cpp b/src/ATen/native/sparse/xpu/SparseCsrTensorMath.cpp
@@ -1,9 +1,12 @@
+#include <ATen/native/sparse/SparseStubs.h>
 #include <ATen/native/sparse/xpu/sycl/SparseCsrTensorMathKernels.h>
 #include <xpu/ATen/ops/_convert_indices_from_coo_to_csr_native.h>
 #include <xpu/ATen/ops/_convert_indices_from_csr_to_coo_native.h>
 
 namespace at::native {
 
+using namespace at::sparse;
+
 TORCH_IMPL_FUNC(_convert_indices_from_coo_to_csr_structured_xpu)
 (const Tensor& input,
  const int64_t size,
@@ -23,4 +26,21 @@ TORCH_IMPL_FUNC(_convert_indices_from_csr_to_coo_structured_xpu)
       crow_indices, col_indices, out_int32, transpose, result);
 };
 
+Tensor _sparse_csr_sum_xpu(
+    const Tensor& input,
+    IntArrayRef dims_to_sum,
+    bool keepdim,
+    std::optional<ScalarType> dtype) {
+  return xpu::_sparse_csr_sum_xpu_kernel(input, dims_to_sum, keepdim, dtype);
+}
+
+Tensor _sparse_csr_prod_xpu(
+    const Tensor& input,
+    IntArrayRef dims_to_reduce,
+    bool keepdim,
+    std::optional<ScalarType> dtype) {
+  return xpu::_sparse_csr_prod_xpu_kernel(
+      input, dims_to_reduce, keepdim, dtype);
+}
+
 } // namespace at::native