Update

[ghstack-poisoned]

Author: swolchok

kernels/portable/cpu/op_glu.cpp

@@ -155,12 +155,10 @@ Tensor& glu_out(
   const size_t non_negative_dim = dim < 0 ? dim + self.dim() : dim;
   const auto in_dtype = self.scalar_type();
 
-  ET_SWITCH_FLOAT_TYPES(in_dtype, ctx, "glu", CTYPE_IN, [&]() {
-    if (out.scalar_type() == ScalarType::Float) {
-      glu_out_tensor<CTYPE_IN, float>(self, non_negative_dim, out);
-    } else {
-      glu_out_tensor<CTYPE_IN, double>(self, non_negative_dim, out);
-    }
+  ET_SWITCH_FLOATHBF16_TYPES(in_dtype, ctx, "glu", CTYPE_IN, [&]() {
+    ET_SWITCH_FLOATHBF16_TYPES(out.scalar_type(), ctx, "glu", CTYPE_OUT, [&]() {
+      glu_out_tensor<CTYPE_IN, CTYPE_OUT>(self, non_negative_dim, out);
+    });
   });
 
   return out;

kernels/test/op_glu_test.cpp

@@ -128,6 +128,20 @@ TEST_F(OpGluOutTest, AllInputDoubleOutputSupport) {
 #undef TEST_ENTRY
 }
 
+TEST_F(OpGluOutTest, AllInputHalfOutputSupport) {
+#define TEST_ENTRY(ctype, dtype) \
+  test_glu_out<ScalarType::dtype, ScalarType::Half>();
+  ET_FORALL_FLOAT_TYPES(TEST_ENTRY);
+#undef TEST_ENTRY
+}
+
+TEST_F(OpGluOutTest, AllInputBFloat16OutputSupport) {
+#define TEST_ENTRY(ctype, dtype) \
+  test_glu_out<ScalarType::dtype, ScalarType::BFloat16>();
+  ET_FORALL_FLOAT_TYPES(TEST_ENTRY);
+#undef TEST_ENTRY
+}
+
 TEST_F(OpGluOutTest, InfinityAndNANTest) {
   TensorFactory<ScalarType::Float> tf;
   const std::vector<int32_t> sizes = {4, 2};