Fix & cleanup op logit (#695)

Summary:
Pull Request resolved: #695

ghstack-source-id: 203341583
exported-using-ghexport

Reviewed By: SS-JIA

Differential Revision: D49625518

fbshipit-source-id: 1f5922e0f3ca7d2f36f69eaa07e9bc42618ca93d

Author: manuelcandales

kernels/portable/cpu/op_logit.cpp

@@ -25,12 +25,13 @@ Tensor& logit_out(
   (void)ctx;
 
   // Resize for dynamic shape
-  auto error = resize_tensor(out, in.sizes());
-  ET_CHECK_MSG(error == Error::Ok, "Failed to resize output tensor.");
-  ET_CHECK_SAME_SHAPE2(in, out);
+  ET_KERNEL_CHECK(
+      ctx, resize_tensor(out, in.sizes()) == Error::Ok, InvalidArgument, out);
 
-  ET_SWITCH_REAL_TYPES_AND(Bool, in.scalar_type(), ctx, "logit", CTYPE_IN, [&] {
-    ET_SWITCH_FLOAT_TYPES(out.scalar_type(), ctx, "logit", CTYPE_OUT, [&] {
+  ScalarType in_type = in.scalar_type();
+  ScalarType out_type = out.scalar_type();
+  ET_SWITCH_REAL_TYPES_AND(Bool, in_type, ctx, __func__, CTYPE_IN, [&] {
+    ET_SWITCH_FLOAT_TYPES(out_type, ctx, __func__, CTYPE_OUT, [&] {
       apply_unary_map_fn(
           [eps](const CTYPE_IN val_in) {
             CTYPE_OUT xi = static_cast<CTYPE_OUT>(val_in);
@@ -41,7 +42,6 @@ Tensor& logit_out(
                 xi = 1 - eps.value();
               }
             }
-            ET_CHECK_MSG(xi > 0.0 && xi < 1.0, "input must be in (0, 1).");
             return static_cast<CTYPE_OUT>(
                 log(xi / (static_cast<CTYPE_OUT>(1.0) - xi)));
           },

kernels/test/op_logit_test.cpp

@@ -37,8 +37,10 @@ void test_integer_logit_out() {
   // Destination for the logit operator.
   Tensor out = tf_out.zeros(sizes);
 
-  ET_EXPECT_KERNEL_FAILURE(
-      op_logit_out(tf.make(sizes, /*data=*/{1, 2, 4, 8}), 0, out));
+  op_logit_out(tf.make(sizes, /*data=*/{1, 2, 4, 8}), 0, out);
+  EXPECT_TENSOR_CLOSE(
+      out,
+      tf_out.make(sizes, /*data=*/{INFINITY, INFINITY, INFINITY, INFINITY}));
 }
 
 template <>
@@ -79,19 +81,13 @@ void test_integer_logit_out_eps_set() {
 }
 
 TEST(OpLogitOutKernelTest, AllRealInputFloatOutputSupport) {
-  if (torch::executor::testing::SupportedFeatures::get()->is_aten) {
-    GTEST_SKIP() << "ATen kernel can handle this";
-  }
 #define TEST_ENTRY(ctype, dtype) \
   test_integer_logit_out<ScalarType::dtype, ScalarType::Float>();
   ET_FORALL_REAL_TYPES(TEST_ENTRY);
 #undef TEST_ENTRY
 }
 
 TEST(OpLogitOutKernelTest, AllRealInputDoubleOutputSupport) {
-  if (torch::executor::testing::SupportedFeatures::get()->is_aten) {
-    GTEST_SKIP() << "ATen kernel can handle this";
-  }
 #define TEST_ENTRY(ctype, dtype) \
   test_integer_logit_out<ScalarType::dtype, ScalarType::Double>();
   ET_FORALL_REAL_TYPES(TEST_ENTRY);