fix scalar tensor.mean

oneflow/core/common/wrap_dim_utils.h

@@ -50,8 +50,8 @@ static inline Maybe<std::bitset<dim_bitset_size>> dim_list_to_bitset(
   std::bitset<dim_bitset_size> seen;
   for (int32_t i = 0; i < dims.size(); i++) {
     size_t dim = JUST(maybe_wrap_dim(dims[i], ndims));
-    CHECK_OR_RETURN(!seen[dim]) << Error::RuntimeError() << "The dim " << dim
-                                << " appears multiple times in the list of dims";
+    CHECK_OR_RETURN_ERROR(!seen[dim]) << Error::RuntimeError() << "The dim " << dim
+                                      << " appears multiple times in the list of dims";
     seen[dim] = true;
   }
   return seen;

oneflow/core/functional/impl/common.cpp

@@ -35,23 +35,18 @@ bool IsScalarTensor(const std::shared_ptr<Tensor>& x) {
 
 Maybe<std::vector<int32_t>> CheckAxis(const std::vector<int32_t>& axis, const int32_t& ndim) {
   const int32_t naxis = axis.size();
-
+  int32_t reduce_ndim = naxis;
+  if (naxis == 0 || ndim == 0) { reduce_ndim = ndim; };
+  std::vector<int32_t> reduce_axis(reduce_ndim);
   if (naxis == 0) {
-    std::vector<int32_t> reduce_axis(ndim);
     std::iota(reduce_axis.begin(), reduce_axis.end(), 0);
-    return reduce_axis;
   } else {
-    std::vector<int32_t> reduce_axis(naxis);
-    std::vector<int32_t> axis_num(ndim);
+    JUST(dim_list_to_bitset(axis, ndim));  // checking axis[dim]'s validation
     for (int32_t i = 0; i < naxis; i++) {
-      reduce_axis[i] = JUST(maybe_wrap_dim(axis[i], ndim));
-      axis_num[reduce_axis[i]]++;
-      CHECK_OR_RETURN(axis_num[reduce_axis[i]] < 2)
-          << Error::RuntimeError() << "dim " << reduce_axis[i]
-          << " appears multiple times in the list of dims";
+      if (i < reduce_ndim) { reduce_axis[i] = JUST(maybe_wrap_dim(axis[i], ndim)); };
     }
-    return reduce_axis;
   }
+  return reduce_axis;
 }
 
 Maybe<void> CheckInplaceValid(const std::shared_ptr<Tensor>& x) {

python/oneflow/test/modules/test_mean.py

@@ -86,6 +86,13 @@ def test_mean_with_random_data(test_case):
         x = random_tensor(ndim=4, dtype=float).to(device)
         return torch.mean(x, dim)
 
+    @autotest(n=5)
+    def test_mean_with_scalar_data(test_case):
+        device = random_device()
+        x = random_tensor(ndim=4, dtype=float).to(device).mean()
+        y = x.mean(-1)
+        return y
+
 
 if __name__ == "__main__":
     unittest.main()