ENH: Let numpy.size accept multiple axes. (numpy#29240)

* Let numpy.size accept multiple axes.

* Apply suggestions from code review

---------

Co-authored-by: Sebastian Berg <[email protected]>

Author: carlosgmartin

doc/release/upcoming_changes/29240.new_feature.rst

@@ -0,0 +1 @@
+* Let ``np.size`` accept multiple axes.

numpy/_core/fromnumeric.py

@@ -2,6 +2,7 @@
 
 """
 import functools
+import math
 import types
 import warnings
 
@@ -3569,10 +3570,13 @@ 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.
 
+        .. versionchanged:: 2.4
+           Extended to accept multiple axes.
+
     Returns
     -------
     element_count : int
@@ -3590,10 +3594,12 @@ def size(a, axis=None):
     >>> a = np.array([[1,2,3],[4,5,6]])
     >>> np.size(a)
     6
-    >>> np.size(a,1)
+    >>> np.size(a,axis=1)
     3
-    >>> np.size(a,0)
+    >>> np.size(a,axis=0)
     2
+    >>> np.size(a,axis=(0,1))
+    6
 
     """
     if axis is None:
@@ -3602,10 +3608,10 @@ def size(a, axis=None):
         except AttributeError:
             return asarray(a).size
     else:
-        try:
-            return a.shape[axis]
-        except AttributeError:
-            return asarray(a).shape[axis]
+        _shape = shape(a)
+        from .numeric import normalize_axis_tuple
+        axis = normalize_axis_tuple(axis, len(_shape), allow_duplicate=False)
+        return math.prod(_shape[ax] for ax in axis)
 
 
 def _round_dispatcher(a, decimals=None, out=None):

numpy/_core/fromnumeric.pyi

@@ -1397,7 +1397,7 @@ def cumulative_prod(
 
 def ndim(a: ArrayLike) -> int: ...
 
-def size(a: ArrayLike, axis: int | None = ...) -> int: ...
+def size(a: ArrayLike, axis: int | tuple[int, ...] | None = ...) -> int: ...
 
 @overload
 def around(

numpy/_core/tests/test_numeric.py

@@ -291,6 +291,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]]]