Add aten::trace (#446)

Add aten::trace Co-authored-by: Feng Yuan <[email protected]>
intel · Jul 4, 2024 · 41ac6a3 · 41ac6a3
1 parent 662ab6a
commit 41ac6a3
Show file tree

Hide file tree

Showing 4 changed files with 8 additions and 1 deletion.
diff --git a/src/ATen/native/xpu/TriangluarOps.cpp b/src/ATen/native/xpu/TriangluarOps.cpp
@@ -68,4 +68,10 @@ Tensor& XPUNativeFunctions::triu_(Tensor& self, int64_t diagonal) {
   xpu::check_inplace(self, self.sizes(), self.options());
   return triu_out(self, diagonal, self);
 }
+
+Tensor XPUNativeFunctions::trace(const Tensor& self) {
+  TORCH_CHECK(self.dim() == 2, "expected a matrix");
+  return self.diagonal().sum();
+}
+
 } // namespace at
diff --git a/src/ATen/native/xpu/XPUFallback.template b/src/ATen/native/xpu/XPUFallback.template
@@ -370,7 +370,6 @@ TORCH_LIBRARY_IMPL(aten, XPU, m) {
     "topk.values",
     "_to_sparse",
     "_to_sparse_csr",
-    "trace",
     "triangular_solve.X",
     "tril_indices",
     "triu_indices",

diff --git a/test/xpu/xpu_test_utils.py b/test/xpu/xpu_test_utils.py
@@ -27,6 +27,7 @@
     "view_as_real",
     "view_as_complex",
     "view",
+    "trace",
     "resize_",
     "resize_as_",
     "add",

diff --git a/yaml/xpu_functions.yaml b/yaml/xpu_functions.yaml
@@ -453,6 +453,7 @@ supported:
   - _cdist_forward
   - _pin_memory
   - is_pinned
+  - trace
   - reflection_pad2d
   - reflection_pad2d.out
   - reflection_pad2d_backward