diff --git a/pymnn/pip_package/MNN/numpy/__init__.py b/pymnn/pip_package/MNN/numpy/__init__.py index a9ea5c390c..dfd78c4857 100644 --- a/pymnn/pip_package/MNN/numpy/__init__.py +++ b/pymnn/pip_package/MNN/numpy/__init__.py @@ -1301,8 +1301,8 @@ def quantile(a, q, axis=None, out=None, overwrite_input=False, interpolation='li def median(a, axis=None, out=None, overwrite_input=False, keepdims=False):__not_impl() def mean(a, axis=None, dtype=float32, out=None, keepdims=False): a = _F.cast(a, dtype) - if axis is None: return asscalar(_F.reduce_mean(ravel(a))) - return _F.reduce_mean(a, axis, keepdims) + if axis is None and not keepdims: return asscalar(_F.reduce_mean(ravel(a))) + return _F.reduce_mean(a, [] if axis is None else axis, keepdims) def average(a, axis=None, weights=None, returned=False): if weights is not None: raise ValueError('MNN.numpy average not support `weights`.') diff --git a/pymnn/test/unit_test.py b/pymnn/test/unit_test.py index 3f118362ff..71132c74b1 100644 --- a/pymnn/test/unit_test.py +++ b/pymnn/test/unit_test.py @@ -1187,6 +1187,19 @@ def test_searching_counting(self): self.assertEqualVar(mp.flatnonzero(x), np.flatnonzero(x_)) self.assertEqualVar(mp.count_nonzero(x, 0), np.count_nonzero(x_, 0)) def test_statistics(self): + for shape in [(3,), (2, 3), (2, 3, 2), (2, 1, 3, 2)]: + data = (np.arange(np.prod(shape), dtype=np.float32).reshape(shape) % 7 - 3) * 0.5 + for axis in [None] + list(range(len(shape))): + for keepdims in [False, True]: + with self.subTest(shape=shape, axis=axis, keepdims=keepdims): + value = mp.array(data) + actual = mp.mean(value, axis=axis, keepdims=keepdims) + expected = np.mean(data, axis=axis, keepdims=keepdims) + if keepdims or axis is not None: + self.assertEqualVar(actual, expected) + else: + self.assertAlmostEqual(actual, float(expected), places=5) + self.assertEqualVar(value, data) x = mp.array([[1,2,3], [4,6,5]]) x_ = np.array([[1,2,3], [4,6,5]]) self.assertEqual(mp.amin(x), np.amin(x_))