Merge branch 'main' into export-D72185131

Author: kkolur76

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;
 

tools/cmake/Codegen.cmake

@@ -83,19 +83,19 @@ function(generate_bindings_for_kernels)
   execute_process(
     COMMAND
       "${PYTHON_EXECUTABLE}" -c
-      "from distutils.sysconfig import get_python_lib;print(get_python_lib())"
-    OUTPUT_VARIABLE site-packages-out
-    ERROR_VARIABLE site-packages-out-error
-    RESULT_VARIABLE site-packages-result
+      "import torchgen;import os; print(os.path.dirname(torchgen.__file__))"
+    OUTPUT_VARIABLE torchgen-out
+    ERROR_VARIABLE torchgen-out-error
+    RESULT_VARIABLE torchgen-result
     WORKING_DIRECTORY ${CMAKE_CURRENT_BINARY_DIR}
     OUTPUT_STRIP_TRAILING_WHITESPACE
   )
-  file(GLOB_RECURSE _torchgen_srcs "${site-packages-out}/torchgen/*.py")
+  file(GLOB_RECURSE _torchgen_srcs "${torchgen-out}/*.py")
   set(_gen_command
       "${PYTHON_EXECUTABLE}" -m torchgen.gen_executorch
       --source-path=${EXECUTORCH_ROOT}/codegen --install-dir=${_out_dir}
-      --tags-path=${site-packages-out}/torchgen/packaged/ATen/native/tags.yaml
-      --aten-yaml-path=${site-packages-out}/torchgen/packaged/ATen/native/native_functions.yaml
+      --tags-path=${torchgen-out}/packaged/ATen/native/tags.yaml
+      --aten-yaml-path=${torchgen-out}/packaged/ATen/native/native_functions.yaml
       --op-selection-yaml-path=${_oplist_yaml}
     )
   if(GEN_ADD_EXCEPTION_BOUNDARY)