Fix bug in elu

Differential Revision: D81359053

Pull Request resolved: #13829

Author: manuelcandales

kernels/optimized/cpu/op_elu.cpp

@@ -28,12 +28,9 @@ void elu(
   CTYPE* out_data = out.mutable_data_ptr<CTYPE>();
   using MathT =
       std::conditional_t<c10::is_reduced_floating_point_v<CTYPE>, float, CTYPE>;
-  MathT math_alpha = 0;
-  MathT math_scale = 0;
-  MathT math_input_scale = 0;
-  ET_EXTRACT_SCALAR(alpha, math_alpha);
-  ET_EXTRACT_SCALAR(scale, math_scale);
-  ET_EXTRACT_SCALAR(input_scale, math_input_scale);
+  const auto math_alpha = utils::scalar_to<MathT>(alpha);
+  const auto math_scale = utils::scalar_to<MathT>(scale);
+  const auto math_input_scale = utils::scalar_to<MathT>(input_scale);
   const auto scalar_func =
       at::native::get_scalar_elu_elementwise_func<CTYPE, MathT>(
           math_alpha, math_scale, math_input_scale);

kernels/portable/cpu/op_elu.cpp

@@ -37,12 +37,9 @@ Tensor& elu_out(
   ET_SWITCH_FLOATHBF16_TYPES(in.scalar_type(), ctx, op_name, CTYPE, [&]() {
     using MathT = std::
         conditional_t<c10::is_reduced_floating_point_v<CTYPE>, float, CTYPE>;
-    MathT math_alpha = 0;
-    MathT math_scale = 0;
-    MathT math_input_scale = 0;
-    ET_EXTRACT_SCALAR(alpha, math_alpha);
-    ET_EXTRACT_SCALAR(scale, math_scale);
-    ET_EXTRACT_SCALAR(input_scale, math_input_scale);
+    const auto math_alpha = utils::scalar_to<MathT>(alpha);
+    const auto math_scale = utils::scalar_to<MathT>(scale);
+    const auto math_input_scale = utils::scalar_to<MathT>(input_scale);
     const auto negcoef = math_alpha * math_scale;
     utils::apply_unitensor_elementwise_fn<
         CTYPE,

kernels/test/op_elu_test.cpp

@@ -43,7 +43,7 @@ class OpEluTest : public OperatorTest {
 
     Tensor out = tf.zeros(sizes);
 
-    // Run full gelu.
+    // Run full elu.
     op_elu_out(in, 1.25, 1, 1, out);
 
     // Check that it matches the expected output.
@@ -93,3 +93,18 @@ TEST_F(OpEluTest, MismatchedOutputDtypeDies) {
 
   ET_EXPECT_KERNEL_FAILURE(context_, op_elu_out(a, 1, 1, 1, out));
 }
+
+TEST_F(OpEluTest, MixedScalarTypes) {
+  TensorFactory<ScalarType::Float> tf_float;
+
+  const std::vector<int32_t> sizes = {2, 2};
+
+  Tensor in = tf_float.ones(sizes);
+  Tensor out = tf_float.zeros(sizes);
+
+  op_elu_out(in, true, 1.0, 1.0, out);
+  EXPECT_TENSOR_CLOSE(out, tf_float.ones(sizes));
+
+  op_elu_out(in, false, true, 3, out);
+  EXPECT_TENSOR_CLOSE(out, tf_float.ones(sizes));
+}