ENH: Let numpy.size accept multiple axes.

carlosgmartin · carlosgmartin · commit 8f17ccf9b537 · 2025-06-19T19:00:08.000-04:00
diff --git a/numpy/_core/fromnumeric.py b/numpy/_core/fromnumeric.py
@@ -2,6 +2,7 @@
 
 """
 import functools
+import math
 import types
 import warnings
 
@@ -3572,8 +3573,8 @@ def size(a, axis=None):
     ----------
     a : array_like
         Input data.
-    axis : int, optional
-        Axis along which the elements are counted.  By default, give
+    axis : None or int or tuple of ints, optional
+        Axis or axes along which the elements are counted.  By default, give
         the total number of elements.
 
     Returns
@@ -3604,6 +3605,12 @@ def size(a, axis=None):
             return a.size
         except AttributeError:
             return asarray(a).size
+    elif isinstance(axis, tuple):
+        try:
+            shape = a.shape
+        except AttributeError:
+            shape = asarray(a).shape
+        return math.prod(shape[ax] for ax in axis)
     else:
         try:
             return a.shape[axis]
diff --git a/numpy/_core/tests/test_numeric.py b/numpy/_core/tests/test_numeric.py
@@ -292,6 +292,10 @@ def test_size(self):
         assert_(np.size(A) == 6)
         assert_(np.size(A, 0) == 2)
         assert_(np.size(A, 1) == 3)
+        assert_(np.size(A, () == 1))
+        assert_(np.size(A, (0,) == 2))
+        assert_(np.size(A, (1,) == 3))
+        assert_(np.size(A, (0, 1) == 6))
 
     def test_squeeze(self):
         A = [[[1, 1, 1], [2, 2, 2], [3, 3, 3]]]
