Repository: mars-project/mars

Issue:
[BUG] SUM on axis=-1  return wrong shape
<!--
Thank you for your contribution!

Please review https://github.com/mars-project/mars/blob/master/CONTRIBUTING.rst before opening an issue.
-->

**Describe the bug**

SUM on axis=-1  return wrong shape.

**To Reproduce**

```
In [228]: X = mt.random.rand(3, 4, 5)                                                                                                                                                                       

In [229]: X.sum(axis=-1).shape                                                                                                                                                                              
Out[229]: (3, 4, 5)  # wrong

In [230]: X.sum(axis=2).shape                                                                                                                                                                               
Out[230]: (3, 4)  # right
```


Discussion:

Resolving patch:
diff --git a/mars/tensor/expressions/reduction/core.py b/mars/tensor/expressions/reduction/core.py
--- a/mars/tensor/expressions/reduction/core.py
+++ b/mars/tensor/expressions/reduction/core.py
@@ -65,7 +65,7 @@ def _call(self, a, out):
         else:
             axis = lrange(len(a.shape)) if axis is None else axis
             if not isinstance(axis, Iterable):
-                axis = (axis,)
+                axis = (validate_axis(a.ndim, axis),)
             axis = set(axis)
 
             shape = tuple(s if i not in axis else 1 for i, s in enumerate(a.shape)
@@ -238,10 +238,7 @@ def _get_arg_axis(axis, ndim):
             axis = tuple(range(ndim))
             ravel = True
         elif isinstance(axis, six.integer_types):
-            if axis < 0:
-                axis += ndim
-            if axis < 0 or axis >= ndim:
-                raise ValueError("axis entry is out of bounds")
+            axis = validate_axis(ndim, axis)
             axis = (axis,)
             ravel = ndim == 1
         else:
diff --git a/mars/tensor/expressions/utils.py b/mars/tensor/expressions/utils.py
--- a/mars/tensor/expressions/utils.py
+++ b/mars/tensor/expressions/utils.py
@@ -96,8 +96,7 @@ def random_state_data(n, random_state=None):
 
 def validate_axis(ndim, axis):
     if axis >= ndim or axis < -ndim:
-        raise ValueError('Axis must be between -%d and %d, got %d' %
-                         (ndim, ndim - 1, axis))
+        raise np.AxisError(axis, ndim=ndim)
 
     return axis if axis >= 0 else ndim + axis
 


Test patch:
diff --git a/mars/tensor/expressions/tests/test_reduction.py b/mars/tensor/expressions/tests/test_reduction.py
--- a/mars/tensor/expressions/tests/test_reduction.py
+++ b/mars/tensor/expressions/tests/test_reduction.py
@@ -16,6 +16,8 @@
 
 import unittest
 
+import numpy as np
+
 from mars.tensor.expressions.datasource import ones, tensor
 from mars.operands import MeanCombine, MeanChunk, Mean, Concatenate, Argmax, Argmin,\
     ArgmaxChunk, ArgmaxCombine, ArgminChunk, ArgminCombine
@@ -50,6 +52,16 @@ def testBaseReduction(self):
             self.assertEqual(res.shape, (10,))
             self.assertEqual(calc_shape(res), res.shape)
 
+            with self.assertRaises(np.AxisError):
+                f(ones((10, 8), chunk_size=3), axis=2)
+
+            res = f(ones((10, 8), chunk_size=3), axis=-1)
+            self.assertEqual(res.shape, (10,))
+            self.assertEqual(calc_shape(res), res.shape)
+
+            with self.assertRaises(np.AxisError):
+                f(ones((10, 8), chunk_size=3), axis=-3)
+
             res = f(ones((10, 8), chunk_size=3), keepdims=True)
             self.assertEqual(res.shape, (1, 1))
             self.assertEqual(calc_shape(res), res.shape)
@@ -100,6 +112,16 @@ def testMeanReduction(self):
         self.assertEqual(res.shape, (10,))
         self.assertEqual(calc_shape(res), res.shape)
 
+        with self.assertRaises(np.AxisError):
+            mean(ones((10, 8), chunk_size=3), axis=2)
+
+        res = mean(ones((10, 8), chunk_size=3), axis=-1)
+        self.assertEqual(res.shape, (10,))
+        self.assertEqual(calc_shape(res), res.shape)
+
+        with self.assertRaises(np.AxisError):
+            mean(ones((10, 8), chunk_size=3), axis=-3)
+
         res = mean(ones((10, 8), chunk_size=3), keepdims=True)
         self.assertEqual(res.shape, (1, 1))
         self.assertEqual(calc_shape(res), res.shape)
@@ -153,6 +175,8 @@ def testArgReduction(self):
 
         self.assertRaises(TypeError, lambda: argmax(ones((10, 8, 10), chunk_size=3), axis=(0, 1)))
         self.assertRaises(TypeError, lambda: argmin(ones((10, 8, 10), chunk_size=3), axis=(0, 1)))
+        self.assertRaises(np.AxisError, lambda: argmin(ones((10, 8, 10), chunk_size=3), axis=3))
+        self.assertRaises(np.AxisError, lambda: argmin(ones((10, 8, 10), chunk_size=3), axis=-4))
 
     def testCumReduction(self):
         cumsum = lambda x, *args, **kwargs: x.cumsum(*args, **kwargs).tiles()
@@ -178,6 +202,20 @@ def testCumReduction(self):
         self.assertEqual(calc_shape(res2), res2.shape)
         self.assertEqual(calc_shape(res2.chunks[0]), res2.chunks[0].shape)
 
+        res1 = cumsum(ones((10, 8, 8), chunk_size=3), axis=-2)
+        res2 = cumprod(ones((10, 8, 8), chunk_size=3), axis=-2)
+        self.assertEqual(res1.shape, (10, 8, 8))
+        self.assertEqual(calc_shape(res1), res1.shape)
+        self.assertEqual(calc_shape(res1.chunks[0]), res1.chunks[0].shape)
+        self.assertEqual(res2.shape, (10, 8, 8))
+        self.assertEqual(calc_shape(res2), res2.shape)
+        self.assertEqual(calc_shape(res2.chunks[0]), res2.chunks[0].shape)
+
+        with self.assertRaises(np.AxisError):
+            cumsum(ones((10, 8), chunk_size=3), axis=2)
+        with self.assertRaises(np.AxisError):
+            cumsum(ones((10, 8), chunk_size=3), axis=-3)
+
     def testAllReduction(self):
         o = tensor([False])