diff options
Diffstat (limited to 'numpy/lib/tests')
| -rw-r--r-- | numpy/lib/tests/test_function_base.py | 12 |
1 files changed, 12 insertions, 0 deletions
diff --git a/numpy/lib/tests/test_function_base.py b/numpy/lib/tests/test_function_base.py index 64318255b..9cc18a5e4 100644 --- a/numpy/lib/tests/test_function_base.py +++ b/numpy/lib/tests/test_function_base.py @@ -360,6 +360,18 @@ class TestAverage: assert_(np.average(y3, weights=w3).dtype == np.result_type(y3, w3)) + # test weights with `keepdims=False` and `keepdims=True` + x = np.array([2, 3, 4]).reshape(3, 1) + w = np.array([4, 5, 6]).reshape(3, 1) + + actual = np.average(x, weights=w, axis=1, keepdims=False) + desired = np.array([2., 3., 4.]) + assert_array_equal(actual, desired) + + actual = np.average(x, weights=w, axis=1, keepdims=True) + desired = np.array([[2.], [3.], [4.]]) + assert_array_equal(actual, desired) + def test_returned(self): y = np.array([[1, 2, 3], [4, 5, 6]]) |