Fix more shape tests for 0-D tensor

Summary: Special handling ops for 0-D tensor

Reviewed By: manuelcandales

Differential Revision: D48323627

fbshipit-source-id: 06b3bdfd55f9a5e8569168875b902402b1f96c66

Author: kirklandsign

kernels/portable/cpu/op_amax.cpp

@@ -28,8 +28,10 @@ void check_preconditions(
     Tensor& out) {
   ET_CHECK_SAME_DTYPE2(in, out);
   check_dim_list_is_valid(in, dim_list);
-  for (const auto& d : dim_list) {
-    ET_CHECK_NON_ZERO_DIM_SIZE(d, in);
+  if (in.dim() != 0) {
+    for (const auto& d : dim_list) {
+      ET_CHECK_NON_ZERO_DIM_SIZE(d, in);
+    }
   }
   ET_CHECK_MSG(
       out.dim() == compute_reduced_out_dim(in, dim_list, keepdim),

kernels/portable/cpu/op_amin.cpp

@@ -28,8 +28,10 @@ void check_preconditions(
     Tensor& out) {
   ET_CHECK_SAME_DTYPE2(in, out);
   check_dim_list_is_valid(in, dim_list);
-  for (const auto& d : dim_list) {
-    ET_CHECK_NON_ZERO_DIM_SIZE(d, in);
+  if (in.dim() != 0) {
+    for (const auto& d : dim_list) {
+      ET_CHECK_NON_ZERO_DIM_SIZE(d, in);
+    }
   }
   ET_CHECK_MSG(
       out.dim() == compute_reduced_out_dim(in, dim_list, keepdim),

kernels/portable/cpu/op_argmax.cpp

@@ -27,6 +27,12 @@ void check_preconditions(
     optional<int64_t> dim,
     bool keepdim,
     Tensor& out) {
+  if (in.dim() == 0) {
+    if (dim.has_value()) {
+      ET_CHECK(dim.value() == 0 || dim.value() == -1);
+    }
+    return;
+  }
   if (dim.has_value()) {
     ET_CHECK_VALID_DIM(dim.value(), in.dim());
     ET_CHECK_NON_ZERO_DIM_SIZE(dim.value(), in);

kernels/portable/cpu/op_argmin.cpp

@@ -27,6 +27,12 @@ void check_preconditions(
     optional<int64_t> dim,
     bool keepdim,
     Tensor& out) {
+  if (in.dim() == 0) {
+    if (dim.has_value()) {
+      ET_CHECK(dim.value() == 0 || dim.value() == -1);
+    }
+    return;
+  }
   if (dim.has_value()) {
     ET_CHECK_VALID_DIM(dim.value(), in.dim());
     ET_CHECK_NON_ZERO_DIM_SIZE(dim.value(), in);

kernels/portable/cpu/op_index_select.cpp

@@ -32,11 +32,15 @@ void check_index_select_args(
     const Tensor& index,
     Tensor& output) {
   // Check dim. The dim planed to be selected on shall exist in input
-  ET_CHECK_MSG(
-      dim >= 0 && dim < input.dim(),
-      "dim %" PRId64 " out of range [0,%zd)",
-      dim,
-      input.dim());
+  if (input.dim() == 0) {
+    ET_CHECK(dim == 0);
+  } else {
+    ET_CHECK_MSG(
+        dim >= 0 && dim < input.dim(),
+        "dim %" PRId64 " out of range [0,%zd)",
+        dim,
+        input.dim());
+  }
 
   // Input output should have the same dim
   ET_CHECK_MSG(
@@ -62,7 +66,10 @@ void check_index_select_args(
   // Index should be a 1-D LongTensor, check if any index is out of bound
   ET_CHECK_MSG(
       index.scalar_type() == ScalarType::Long, "index scalar_type not long");
-  ET_CHECK_MSG(index.dim() == 1, "index.dim() %zd != 1", index.dim());
+  ET_CHECK_MSG(
+      index.dim() == 1 || index.dim() == 0,
+      "index.dim() %zd != 1 or 0",
+      index.dim());
 
   const int64_t* src = index.mutable_data_ptr<int64_t>();
   for (auto i = 1; i < index.numel(); i++) {

kernels/portable/cpu/op_log_softmax.cpp

@@ -40,7 +40,11 @@ void check_preconditions(
   // Ensure in has value
   ET_CHECK_MSG(in.numel() > 0, "in.numel() %zd <= 0", in.numel());
   // Ensure dim is valid
-  ET_CHECK_VALID_DIM(dim, in.dim());
+  if (in.dim() == 0) {
+    ET_CHECK_MSG(dim == 0 || dim == -1, "dim must be 0 or -1 for 0-D tensor");
+  } else {
+    ET_CHECK_VALID_DIM(dim, in.dim());
+  }
   ET_CHECK_DEFAULT_OR_CHANNELSLAST_DIMORDER(in);
   ET_CHECK_DEFAULT_OR_CHANNELSLAST_DIMORDER(out);
 }

kernels/portable/cpu/op_select_copy.cpp

@@ -27,32 +27,38 @@ void check_and_update_select_copy_int_out_args(
     int64_t dim,
     int64_t index,
     Tensor output) {
-  // Support python-style negative indexing. E.g., for the shape {2, 3, 4},
-  // dim = -1 would refer to dim[2], dim = -2 would refer to dim[1], and so on.
-
-  // The dim planed to be selected on shall exist in input
-  ET_CHECK_MSG(
-      dim >= -input.dim() && dim < input.dim(),
-      "dim %" PRId64 " out of range [-%zd,%zd)",
-      dim,
-      input.dim(),
-      input.dim());
-
-  // The index shall be valid in the given dimenson
-  ET_CHECK_MSG(
-      index >= -input.size(dim) && index < input.size(dim),
-      "index %" PRId64 " out of range [-%zd,%zd) at input.size( %" PRId64 ")",
-      index,
-      input.size(dim),
-      input.size(dim),
-      dim);
+  if (input.dim() == 0) {
+    ET_CHECK(dim == 0 || dim == -1);
+  } else {
+    // Support python-style negative indexing. E.g., for the shape {2, 3, 4},
+    // dim = -1 would refer to dim[2], dim = -2 would refer to dim[1], and so
+    // on.
+
+    // The dim planed to be selected on shall exist in input
+    ET_CHECK_MSG(
+        dim >= -input.dim() && dim < input.dim(),
+        "dim %" PRId64 " out of range [-%zd,%zd)",
+        dim,
+        input.dim(),
+        input.dim());
+
+    // Support python-style negative indexing
+    if (dim < 0) {
+      dim += input.dim();
+    }
 
-  // Support python-style negative indexing
-  if (dim < 0) {
-    dim += input.dim();
-  }
-  if (index < 0) {
-    index += input.size(dim);
+    // The index shall be valid in the given dimenson
+    ET_CHECK_MSG(
+        index >= -input.size(dim) && index < input.size(dim),
+        "index %" PRId64 " out of range [-%zd,%zd) at input.size( %" PRId64 ")",
+        index,
+        input.size(dim),
+        input.size(dim),
+        dim);
+
+    if (index < 0) {
+      index += input.size(dim);
+    }
   }
 
   // Input dtype shall match the output dtype.
@@ -71,7 +77,7 @@ void check_and_update_select_copy_int_out_args(
   // - output.size(i) shall equal to input.size(i) if i < dim,
   // - output.size(i) shall equal to input.size(i+1) if i >= dim
 
-  for (size_t d = 0; d < input.dim() - 1; d++) {
+  for (ssize_t d = 0; d < input.dim() - 1; d++) {
     if (d < dim) {
       ET_CHECK_MSG(
           input.size(d) == output.size(d),

kernels/portable/cpu/op_softmax.cpp

@@ -37,10 +37,12 @@ void check_preconditions(
       "in.dim() %zd!= out.dim() %zd",
       in.dim(),
       out.dim());
-  // Ensure in has value
-  ET_CHECK_MSG(in.numel() > 0, "in.numel() %zd <= 0", in.numel());
   // Ensure dim is valid
-  ET_CHECK_VALID_DIM(dim, in.dim());
+  if (in.dim() == 0) {
+    ET_CHECK(dim == 0 || dim == -1);
+  } else {
+    ET_CHECK_VALID_DIM(dim, in.dim());
+  }
   ET_CHECK_DEFAULT_OR_CHANNELSLAST_DIMORDER(in);
   ET_CHECK_DEFAULT_OR_CHANNELSLAST_DIMORDER(out);
 }

kernels/portable/cpu/op_squeeze_copy.cpp

@@ -82,6 +82,9 @@ Tensor& squeeze_copy_dim_out(
   ET_CHECK_SAME_DTYPE2(self, out);
 
   // A valid dim must be in [-self.dim(), self.dim())
+  if (self.dim() == 0 && dim == -1) {
+    dim = 0;
+  }
   ET_CHECK_MSG(
       (self.dim() == 0 && dim == 0) || (dim >= -self.dim() && dim < self.dim()),
       "dim %" PRId64 " out of range [-%zd,%zd)",

kernels/test/op_log_softmax_test.cpp

@@ -93,24 +93,6 @@ TEST(OpLogSoftmaxOutTest, AllDtypesSupported) {
   // for those types.
 }
 
-TEST(OpLogSoftmaxOutTest, EmptyInputOrEmptyOutTensorDies) {
-  if (SupportedFeatures::get()->is_aten) {
-    GTEST_SKIP() << "ATen currently supports empty input or out";
-  }
-
-  TensorFactory<ScalarType::Float> tff;
-
-  Tensor x = tff.make({2, 2, 0}, {});
-
-  // Make an empty out tensor and demonstrate that it's empty.
-  Tensor out = tff.make({2, 2, 0}, {});
-
-  EXPECT_EQ(out.numel(), 0);
-
-  ET_EXPECT_KERNEL_FAILURE(
-      op_log_softmax_out(x, /*dim=*/1, /*half_to_float*/ false, out));
-}
-
 TEST(OpLogSoftmaxOutTest, MismatchedDimensionsDies) {
   if (SupportedFeatures::get()->is_aten) {
     GTEST_SKIP() << "ATen currently supports mismatched dimensions";