From e67fe9f1a86013dbfd4a5ad2bd69aa6e7f60e1fb Mon Sep 17 00:00:00 2001 From: Srimukh Sripada Date: Mon, 27 Jun 2022 05:06:00 +0200 Subject: 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 --- numpy/lib/function_base.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) (limited to 'numpy/lib/function_base.py') 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") -- cgit v1.2.1