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()