Merge branch 'main' into penghuic/pytest_by_script

intel · Jan 6, 2025 · b7a9117 · b7a9117
2 parents c924f3d + ad8f244
commit b7a9117
Show file tree

Hide file tree

Showing 7 changed files with 934 additions and 4 deletions.
diff --git a/src/ATen/native/xpu/TensorShape.cpp b/src/ATen/native/xpu/TensorShape.cpp
@@ -7,6 +7,7 @@
 #include <ATen/native/TensorShape.h>
 #include <ATen/native/TypeProperties.h>
 #include <ATen/native/xpu/sycl/ShapeKernels.h>
+#include <ATen/native/xpu/sycl/TensorShapeKernels.h>
 #include <comm/RegisterUtils.h>
 #include <xpu/ATen/ops/as_strided_copy_native.h>
 #include <xpu/ATen/ops/as_strided_native.h>
@@ -58,5 +59,25 @@ TORCH_IMPL_FUNC(cat_out_xpu)
       result);
 }
 
+void split_with_sizes_copy_out_xpu(
+    const Tensor& self,
+    IntArrayRef split_sizes,
+    int64_t dim,
+    TensorList out) {
+  xpu::split_with_sizes_copy_out_xpu_kernel(self, split_sizes, dim, out);
+}
+
+Tensor _chunk_cat_xpu(TensorList tensors, int64_t dim, int64_t num_chunks) {
+  return xpu::_chunk_cat_xpu_kernel(tensors, dim, num_chunks);
+}
+
+Tensor& _chunk_cat_out_xpu(
+    TensorList tensors,
+    int64_t dim,
+    int64_t num_chunks,
+    Tensor& out) {
+  return xpu::_chunk_cat_out_xpu_kernel(tensors, dim, num_chunks, out);
+}
+
 } // namespace native
 } // namespace at