Support BFloat16 in flip (#7820)

Partial fix for #7748.
pytorch · Jan 22, 2025 · 39aa22d · 39aa22d
1 parent c86a2f6
commit 39aa22d
Show file tree

Hide file tree

Showing 2 changed files with 18 additions and 12 deletions.
diff --git a/kernels/portable/cpu/op_flip.cpp b/kernels/portable/cpu/op_flip.cpp
@@ -66,7 +66,7 @@ Tensor& flip_out(
 
   constexpr auto name = "flip.out";
 
-  ET_SWITCH_REALHB_TYPES(in.scalar_type(), ctx, name, CTYPE, [&] {
+  ET_SWITCH_REALHBBF16_TYPES(in.scalar_type(), ctx, name, CTYPE, [&] {
     const CTYPE* in_data = in.const_data_ptr<CTYPE>();
     CTYPE* out_data = out.mutable_data_ptr<CTYPE>();
 

diff --git a/kernels/test/op_flip_test.cpp b/kernels/test/op_flip_test.cpp
@@ -33,20 +33,26 @@ class OpFlipOutTest : public ::testing::Test {
     // first.
     torch::executor::runtime_init();
   }
+
+  template <ScalarType DTYPE>
+  void test_1d_dtype() {
+    TensorFactory<DTYPE> tf;
+
+    Tensor input = tf.make({4, 1, 3}, {1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12});
+    int64_t dims_data[1] = {-1};
+    IntArrayRef dims = IntArrayRef(dims_data, 1);
+    Tensor out = tf.zeros({4, 1, 3});
+    Tensor out_expected =
+        tf.make({4, 1, 3}, {3, 2, 1, 6, 5, 4, 9, 8, 7, 12, 11, 10});
+    op_flip_out(input, dims, out);
+    EXPECT_TENSOR_CLOSE(out, out_expected);
+  }
 };
 
 TEST_F(OpFlipOutTest, SmokeTest1Dim) {
-  TensorFactory<ScalarType::Float> tfFloat;
-
-  Tensor input =
-      tfFloat.make({4, 1, 3}, {1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12});
-  int64_t dims_data[1] = {-1};
-  IntArrayRef dims = IntArrayRef(dims_data, 1);
-  Tensor out = tfFloat.zeros({4, 1, 3});
-  Tensor out_expected =
-      tfFloat.make({4, 1, 3}, {3, 2, 1, 6, 5, 4, 9, 8, 7, 12, 11, 10});
-  op_flip_out(input, dims, out);
-  EXPECT_TENSOR_CLOSE(out, out_expected);
+#define TEST_ENTRY(ctype, dtype) test_1d_dtype<ScalarType::dtype>();
+  ET_FORALL_REALHBF16_TYPES(TEST_ENTRY);
+#undef TEST_ENTRY
 }
 
 TEST_F(OpFlipOutTest, SmokeTest2Dims) {