BUG: Ensure summed weights returned by np.average always are correct class (numpy#30522)

mhvk · web-flow · commit 1883d8425000 · 2025-12-26T12:05:46.000-07:00
diff --git a/numpy/lib/_function_base_impl.py b/numpy/lib/_function_base_impl.py
@@ -594,7 +594,7 @@ def average(a, axis=None, weights=None, returned=False, *,
 
     if returned:
         if scl.shape != avg_as_array.shape:
-            scl = np.broadcast_to(scl, avg_as_array.shape).copy()
+            scl = np.broadcast_to(scl, avg_as_array.shape, subok=True).copy()
         return avg, scl
     else:
         return avg
diff --git a/numpy/lib/tests/test_function_base.py b/numpy/lib/tests/test_function_base.py
@@ -497,6 +497,14 @@ class subclass(np.ndarray):
 
         assert_equal(type(np.average(a)), subclass)
         assert_equal(type(np.average(a, weights=w)), subclass)
+        # Ensure a possibly returned sum of weights is correct too.
+        ra, rw = np.average(a, weights=w, returned=True)
+        assert_equal(type(ra), subclass)
+        assert_equal(type(rw), subclass)
+        # Even if it needs to be broadcast.
+        ra, rw = np.average(a, weights=w[0], axis=1, returned=True)
+        assert_equal(type(ra), subclass)
+        assert_equal(type(rw), subclass)
 
     def test_upcasting(self):
         typs = [('i4', 'i4', 'f8'), ('i4', 'f4', 'f8'), ('f4', 'i4', 'f8'),
