BUG: Use `keepdims` during normalization in `np.average` and `np.ma.average` (#21851)

The keepdims flag needs to be applied during the calculation of the sum of
the weights in np.average and np.ma.average.  Not passing it causes weights
to broadcast incorrectly.

Fixes #21850

4 files changed, 23 insertions, 2 deletions

diff --git a/numpy/lib/function_base.py b/numpy/lib/function_base.py
index 843e1b85a..d90c23bfe 100644
--- a/numpy/lib/function_base.py
+++ b/numpy/lib/function_base.py
@@ -542,7 +542,7 @@ def average(a, axis=None, weights=None, returned=False, *,
             wgt = np.broadcast_to(wgt, (a.ndim-1)*(1,) + wgt.shape)
             wgt = wgt.swapaxes(-1, axis)
 
-        scl = wgt.sum(axis=axis, dtype=result_dtype)
+        scl = wgt.sum(axis=axis, dtype=result_dtype, **keepdims_kw)
         if np.any(scl == 0.0):
             raise ZeroDivisionError(
                 "Weights sum to zero, can't be normalized")
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]])
 
diff --git a/numpy/ma/extras.py b/numpy/ma/extras.py
index d90831b9b..d2986012b 100644
--- a/numpy/ma/extras.py
+++ b/numpy/ma/extras.py
@@ -645,7 +645,7 @@ def average(a, axis=None, weights=None, returned=False, *,
             wgt = wgt*(~a.mask)
             wgt.mask |= a.mask
 
-        scl = wgt.sum(axis=axis, dtype=result_dtype)
+        scl = wgt.sum(axis=axis, dtype=result_dtype, **keepdims_kw)
         avg = np.multiply(a, wgt,
                           dtype=result_dtype).sum(axis, **keepdims_kw) / scl
 
diff --git a/numpy/ma/tests/test_extras.py b/numpy/ma/tests/test_extras.py
index 1827edd1f..04bf8cfc2 100644
--- a/numpy/ma/tests/test_extras.py
+++ b/numpy/ma/tests/test_extras.py
@@ -241,6 +241,15 @@ class TestAverage:
         a2dma = average(a2dm, axis=1)
         assert_equal(a2dma, [1.5, 4.0])
 
+    def test_testAverage4(self):
+        # Test that `keepdims` works with average
+        x = np.array([2, 3, 4]).reshape(3, 1)
+        b = np.ma.array(x, mask=[[False], [False], [True]])
+        w = np.array([4, 5, 6]).reshape(3, 1)
+        actual = average(b, weights=w, axis=1, keepdims=True)
+        desired = masked_array([[2.], [3.], [4.]], [[False], [False], [True]])
+        assert_equal(actual, desired)
+
     def test_onintegers_with_mask(self):
         # Test average on integers with mask
         a = average(array([1, 2]))