Special handling for 0-D tensor

Summary:
`ET_CHECK_VALID_DIM` checks whether a tensor has a valid dimension. However, for 0-D tensors, some behaviors are different.

Accessing size/stride for dim=0 is invalid, because the valid dimension is [-0, 0) which is none.
```
torch.tensor(2).size(dim=0)
```

However, some ops allow accessing dim=0 or -1 for a 0-D tensor
```
torch.mean(torch.tensor(2, dtype=float), dim=-1)
```

Therefore, in reduce_util helper functions and ops, we need to special handle that case.

We also want to revisit this check for 0-D tensor case for ops.

Reviewed By: manuelcandales

Differential Revision: D48319644

fbshipit-source-id: 0a394ab4caccbcbb7868ccd371276e9dd047f054

Author: kirklandsign

kernels/portable/cpu/op_max.cpp

@@ -37,8 +37,12 @@ void check_preconditions(
       max_indices.scalar_type() == ScalarType::Long,
       "dtype of the max_indices Tensor expected to be be long.");
   // Ensure dim is valid
-  ET_CHECK_VALID_DIM(dim, in.dim());
-  ET_CHECK_NON_ZERO_DIM_SIZE(dim, in);
+  if (in.dim() == 0) {
+    ET_CHECK(dim == 0 || dim == -1);
+  } else {
+    ET_CHECK_VALID_DIM(dim, in.dim());
+    ET_CHECK_NON_ZERO_DIM_SIZE(dim, in);
+  }
   const auto expected_dim = compute_reduced_out_dim(in, dim, keepdim);
   ET_CHECK_MSG(
       max.dim() == expected_dim && max_indices.dim() == expected_dim,

kernels/portable/cpu/op_min.cpp

@@ -37,8 +37,12 @@ void check_preconditions(
       min_indices.scalar_type() == ScalarType::Long,
       "dtype of the min_indices Tensor expected to be be long.");
   // Ensure dim is valid
-  ET_CHECK_VALID_DIM(dim, in.dim());
-  ET_CHECK_NON_ZERO_DIM_SIZE(dim, in);
+  if (in.dim() == 0) {
+    ET_CHECK(dim == 0 || dim == -1);
+  } else {
+    ET_CHECK_VALID_DIM(dim, in.dim());
+    ET_CHECK_NON_ZERO_DIM_SIZE(dim, in);
+  }
   const auto expected_dim = compute_reduced_out_dim(in, dim, keepdim);
   ET_CHECK_MSG(
       min.dim() == expected_dim && min_indices.dim() == expected_dim,

kernels/portable/cpu/op_transpose_copy.cpp

@@ -32,7 +32,12 @@ void check_preconditions(
     Tensor& out) {
   auto a_dim = a.dim();
   ET_CHECK_MSG(
-      a_dim > 0 && a_dim == out.dim(), "invalid rank of tensor a: %zd", a_dim);
+      a_dim >= 0 && a_dim == out.dim(), "invalid rank of tensor a: %zd", a_dim);
+  if (a_dim == 0) {
+    ET_CHECK(dim0 == 0 || dim0 == -1);
+    ET_CHECK(dim1 == 0 || dim1 == -1);
+    return;
+  }
   ET_CHECK_MSG(
       dim0 >= 0 && dim0 < a_dim,
       "dim0: %" PRId64 " out of bounds [0,%zd)",

kernels/portable/cpu/util/reduce_util.cpp

@@ -222,19 +222,19 @@ size_t compute_reduced_out_size(
   if (dim.has_value()) {
     const auto dim_val = dim.value();
     const size_t non_neg_dim = _normalize_non_neg_d(dim_val, in_dim);
-    for (size_t i = 0; i < non_neg_dim; ++i) {
+    for (ssize_t i = 0; i < non_neg_dim; ++i) {
       sizes_arr[i] = in.size(i);
     }
     if (keepdim) {
       sizes_arr[non_neg_dim] = 1;
-      for (size_t i = non_neg_dim + 1; i < in_dim; ++i) {
+      for (ssize_t i = non_neg_dim + 1; i < in_dim; ++i) {
         sizes_arr[i] = in.size(i);
       }
     } else {
-      for (size_t i = non_neg_dim; i < in_dim - 1; ++i) {
+      for (ssize_t i = non_neg_dim; i < in_dim - 1; ++i) {
         sizes_arr[i] = in.size(i + 1);
       }
-      out_dim = in_dim - 1;
+      out_dim = in_dim == 0 ? 0 : in_dim - 1;
     }
   } else {
     if (keepdim) {

kernels/portable/cpu/util/reduce_util.h

@@ -200,6 +200,8 @@ void apply_over_dim(
   if (in.dim() != 0) {
     ET_CHECK_VALID_DIM(dim.value(), in.dim());
   } else {
+    // Special handling for 0-D tensor; 0 or -1 is valid for PyTorch code
+    // `torch.mean(torch.tensor(2, dtype=float), dim=-1)`
     ET_CHECK(dim.value() == 0 || dim.value() == -1);
     fn(in.numel(), 1, 0);
     return;
@@ -243,7 +245,11 @@ void apply_over_dim(
     const int64_t start = 0,
     const int64_t end = -1) {
   if (dim.has_value()) {
-    ET_CHECK_VALID_DIM(dim.value(), in.dim());
+    if (in.dim() != 0) {
+      ET_CHECK_VALID_DIM(dim.value(), in.dim());
+    } else {
+      ET_CHECK(dim.value() == 0 || dim.value() == -1);
+    }
   }
   ET_CHECK_MSG(
       out_ix < get_out_numel(in, dim),
@@ -255,10 +261,10 @@ void apply_over_dim(
   }
 
   const size_t iter_length = get_reduced_dim_product(in, dim);
-  ET_CHECK_VALID_IX(start, iter_length);
-  ET_CHECK_VALID_IX(end, iter_length);
-  const size_t ustart = ET_NORMALIZE_IX(start, iter_length);
-  const size_t uend = ET_NORMALIZE_IX(end, iter_length);
+  const size_t normalized_start = ET_NORMALIZE_IX(start, iter_length);
+  const size_t normalized_end = ET_NORMALIZE_IX(end, iter_length);
+  const size_t ustart = std::max(normalized_start, size_t(0));
+  const size_t uend = std::min(normalized_end, iter_length - 1);
 
   // If dim is null, iterate over the entire tensor
   if (!dim.has_value()) {
@@ -273,8 +279,12 @@ void apply_over_dim(
   // Compute non-negative dimension value from dim value
   const size_t d = ET_NORMALIZE_IX(dim.value(), in.dim());
 
-  apply_on_flat_and_dim_ix_with_stride_and_base(
-      fn, in.strides()[d], base, ustart, uend);
+  if (in.dim() == 0) {
+    fn(base, ustart);
+  } else {
+    apply_on_flat_and_dim_ix_with_stride_and_base(
+        fn, in.strides()[d], base, ustart, uend);
+  }
 }
 
 /**
@@ -370,7 +380,11 @@ std::tuple<CTYPE_OUT, long> map_reduce_over_dim(
     const exec_aten::optional<int64_t>& dim,
     const size_t out_ix) {
   if (dim.has_value()) {
-    ET_CHECK_VALID_DIM(dim.value(), in.dim());
+    if (in.dim() != 0) {
+      ET_CHECK_VALID_DIM(dim.value(), in.dim());
+    } else {
+      ET_CHECK(dim.value() == 0 || dim.value() == -1);
+    }
   }
 
   ET_CHECK_MSG(

kernels/portable/cpu/util/test/reduce_test.cpp

@@ -438,7 +438,8 @@ TEST(ReduceUtilTest, ApplyOnZeroDimTensorOverDim) {
   TensorFactory<ScalarType::Long> tf;
 
   Tensor in = tf.ones({});
-  ET_EXPECT_DEATH(_apply_over_dim(in, 0), "");
+  _apply_over_dim(in, 0);
+  EXPECT_TENSOR_EQ(in, tf.make({}, {0}));
 }
 
 TEST(ReduceUtilTest, ApplyOnZeroDimTensorOverDimListNull) {