Support Half/BFloat16 in tril (#7893)

swolchok · web-flow · commit c954fbcd54f9 · 2025-01-28T08:24:16.000-08:00
Partial fix for #7748.
diff --git a/kernels/portable/cpu/op_tril.cpp b/kernels/portable/cpu/op_tril.cpp
@@ -158,7 +158,7 @@ Tensor& tril_out(
   clear_out(out);
 
   ScalarType out_type = out.scalar_type();
-  ET_SWITCH_REAL_TYPES_AND(Bool, out_type, ctx, __func__, CTYPE, [&]() {
+  ET_SWITCH_REALHBBF16_TYPES(out_type, ctx, __func__, CTYPE, [&]() {
     tril_kernel<CTYPE>(ctx, self, diagonal, out);
   });
 
diff --git a/kernels/test/op_tril_test.cpp b/kernels/test/op_tril_test.cpp
@@ -727,7 +727,7 @@ class OpTrilTest : public OperatorTest {
     test_tril_out_multi_unequal_dim<ScalarType::DTYPE>(); \
   }
 
-ET_FORALL_REAL_TYPES_AND(Bool, GENERATE_GENERIC_TEST)
+ET_FORALL_REALHBBF16_TYPES(GENERATE_GENERIC_TEST)
 
 // Create generic tests for real dtypes. Tensors have diverse values.
 #define GENERATE_REAL_TEST(_, DTYPE)                          \
@@ -738,7 +738,7 @@ ET_FORALL_REAL_TYPES_AND(Bool, GENERATE_GENERIC_TEST)
     test_tril_out_randint_multi_unequal<ScalarType::DTYPE>(); \
   }
 
-ET_FORALL_REAL_TYPES(GENERATE_REAL_TEST)
+ET_FORALL_REALHBBF16_TYPES(GENERATE_REAL_TEST)
 
 TEST_F(OpTrilTest, InvalidInputShapesDies) {
   TensorFactory<ScalarType::Int> tf;

Original file line number	Diff line number	Diff line change
`@@ -727,7 +727,7 @@ class OpTrilTest : public OperatorTest {`
`727`	`727`	`test_tril_out_multi_unequal_dim<ScalarType::DTYPE>(); \`
`728`	`728`	`}`
`729`	`729`
`730`		`-ET_FORALL_REAL_TYPES_AND(Bool, GENERATE_GENERIC_TEST)`
	`730`	`+ET_FORALL_REALHBBF16_TYPES(GENERATE_GENERIC_TEST)`
`731`	`731`
`732`	`732`	`// Create generic tests for real dtypes. Tensors have diverse values.`
`733`	`733`	`#define GENERATE_REAL_TEST(_, DTYPE) \`
`@@ -738,7 +738,7 @@ ET_FORALL_REAL_TYPES_AND(Bool, GENERATE_GENERIC_TEST)`
`738`	`738`	`test_tril_out_randint_multi_unequal<ScalarType::DTYPE>(); \`
`739`	`739`	`}`
`740`	`740`
`741`		`-ET_FORALL_REAL_TYPES(GENERATE_REAL_TEST)`
	`741`	`+ET_FORALL_REALHBBF16_TYPES(GENERATE_REAL_TEST)`
`742`	`742`
`743`	`743`	`TEST_F(OpTrilTest, InvalidInputShapesDies) {`
`744`	`744`	`TensorFactory<ScalarType::Int> tf;`