Complex Support: diagonal_copy

Differential Revision: D72189747

Pull Request resolved: #9777

Author: manuelcandales

kernels/portable/cpu/op_diagonal_copy.cpp

@@ -100,7 +100,7 @@ Tensor& diagonal_copy_out(
 
   constexpr auto name = "diagonal_copy.out";
 
-  ET_SWITCH_REALHBBF16_TYPES(in.scalar_type(), ctx, name, CTYPE, [&] {
+  ET_SWITCH_ALL_TYPES(in.scalar_type(), ctx, name, CTYPE, [&] {
     diagonal_copy_impl<CTYPE>(in, offset, dim1, dim2, out);
   });
 

kernels/test/op_diagonal_copy_test.cpp

@@ -50,6 +50,38 @@ class OpDiagonalCopyOutTest : public ::testing::Test {
     op_diagonal_copy_out(input, 1, 1, 0, out);
     EXPECT_TENSOR_CLOSE(out, out_expected);
   }
+
+  template <typename CTYPE, ScalarType DTYPE>
+  void run_2d_complex_dtype() {
+    TensorFactory<DTYPE> tf;
+    constexpr auto REAL_DTYPE = executorch::runtime::toRealValueType(DTYPE);
+    using REAL_CTYPE =
+        typename executorch::runtime::ScalarTypeToCppType<REAL_DTYPE>::type;
+    Tensor input = tf.make(
+        {3, 4},
+        {CTYPE{REAL_CTYPE(1), REAL_CTYPE(1)},
+         CTYPE{REAL_CTYPE(2), REAL_CTYPE(2)},
+         CTYPE{REAL_CTYPE(3), REAL_CTYPE(3)},
+         CTYPE{REAL_CTYPE(4), REAL_CTYPE(4)},
+         CTYPE{REAL_CTYPE(5), REAL_CTYPE(5)},
+         CTYPE{REAL_CTYPE(6), REAL_CTYPE(6)},
+         CTYPE{REAL_CTYPE(7), REAL_CTYPE(7)},
+         CTYPE{REAL_CTYPE(8), REAL_CTYPE(8)},
+         CTYPE{REAL_CTYPE(9), REAL_CTYPE(9)},
+         CTYPE{REAL_CTYPE(10), REAL_CTYPE(10)},
+         CTYPE{REAL_CTYPE(11), REAL_CTYPE(11)},
+         CTYPE{REAL_CTYPE(12), REAL_CTYPE(12)}});
+    Tensor out = tf.make(
+        {2},
+        {CTYPE{REAL_CTYPE(0), REAL_CTYPE(0)},
+         CTYPE{REAL_CTYPE(0), REAL_CTYPE(0)}});
+    Tensor out_expected = tf.make(
+        {2},
+        {CTYPE{REAL_CTYPE(5), REAL_CTYPE(5)},
+         CTYPE{REAL_CTYPE(10), REAL_CTYPE(10)}});
+    op_diagonal_copy_out(input, 1, 1, 0, out);
+    EXPECT_TENSOR_CLOSE(out, out_expected);
+  }
 };
 
 TEST_F(OpDiagonalCopyOutTest, SmokeTest2D) {
@@ -58,6 +90,13 @@ TEST_F(OpDiagonalCopyOutTest, SmokeTest2D) {
 #undef TEST_ENTRY
 }
 
+TEST_F(OpDiagonalCopyOutTest, ComplexSmokeTest2D) {
+#define TEST_ENTRY(ctype, dtype) \
+  run_2d_complex_dtype<ctype, ScalarType::dtype>();
+  ET_FORALL_COMPLEXH_TYPES(TEST_ENTRY);
+#undef TEST_ENTRY
+}
+
 TEST_F(OpDiagonalCopyOutTest, SmokeTest3D) {
   TensorFactory<ScalarType::Float> tfFloat;