Update

[ghstack-poisoned]

Author: swolchok

kernels/portable/cpu/op_leaky_relu.cpp

@@ -44,7 +44,7 @@ Tensor& leaky_relu_out(
 
   ET_KERNEL_CHECK(ctx, in_type == out_type, InvalidArgument, out);
 
-  ET_SWITCH_FLOAT_TYPES(in_type, ctx, "leaky_relu.out", CTYPE, [&]() {
+  ET_SWITCH_FLOATHBF16_TYPES(in_type, ctx, "leaky_relu.out", CTYPE, [&]() {
     CTYPE negative_slope_casted;
     ET_SWITCH_SCALAR_OBJ_TYPES(
         sc_type, ctx, "leaky_relu.out", CTYPE_MIN, [&]() {

kernels/test/op_leaky_relu_test.cpp

@@ -29,15 +29,21 @@ class OpLeakyReluTest : public OperatorTest {
     return torch::executor::aten::leaky_relu_outf(
         context_, in, negative_slope, out);
   }
-};
+  template <ScalarType DTYPE>
+  void test_leaky_relu_dtype() {
+    TensorFactory<DTYPE> tf;
+    Tensor in = tf.ones({2, 2});
+    Tensor out = tf.zeros({2, 2});
 
-TEST_F(OpLeakyReluTest, SanityCheck) {
-  TensorFactory<ScalarType::Float> tf;
-  Tensor in = tf.ones({2, 2});
-  Tensor out = tf.zeros({2, 2});
+    Tensor ret = op_leaky_relu_out(in, -0.01, out);
 
-  Tensor ret = op_leaky_relu_out(in, -0.01, out);
+    EXPECT_TENSOR_EQ(out, ret);
+    EXPECT_TENSOR_EQ(out, tf.ones({2, 2}));
+  }
+};
 
-  EXPECT_TENSOR_EQ(out, ret);
-  EXPECT_TENSOR_EQ(out, tf.ones({2, 2}));
+TEST_F(OpLeakyReluTest, SanityCheck) {
+#define TEST_ENTRY(ctype, dtype) test_leaky_relu_dtype<ScalarType::dtype>();
+  ET_FORALL_FLOATHBF16_TYPES(TEST_ENTRY);
+#undef TEST_ENTRY
 }