update dpnp.size to accept tuple of ints for axes

CHANGELOG.md

@@ -30,6 +30,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 * Changed th order of individual FFTs over `axes` for `dpnp.fft.irfftn` to be in forward order [#2524](https://github.com/IntelPython/dpnp/pull/2524)
 * Replaced the use of `numpy.testing.suppress_warnings` with appropriate calls from the warnings module [#2529](https://github.com/IntelPython/dpnp/pull/2529)
 * Improved documentations of `dpnp.ndarray` class and added a page with description of supported constants [#2422](https://github.com/IntelPython/dpnp/pull/2422)
+* Updated `dpnp.size` to accept tuple of ints for `axes` argument [#2536](https://github.com/IntelPython/dpnp/pull/2536)
 
 ### Deprecated
 

dpnp/dpnp_iface_manipulation.py

@@ -46,7 +46,11 @@
 import dpctl
 import dpctl.tensor as dpt
 import numpy
-from dpctl.tensor._numpy_helper import AxisError, normalize_axis_index
+from dpctl.tensor._numpy_helper import (
+    AxisError,
+    normalize_axis_index,
+    normalize_axis_tuple,
+)
 
 import dpnp
 
@@ -3528,8 +3532,8 @@ def size(a, axis=None):
     ----------
     a : array_like
         Input data.
-    axis : {None, int}, optional
-        Axis along which the elements are counted.
+    axis : {None, int, tuple of ints}, optional
+        Axis or axes along which the elements are counted.
         By default, give the total number of elements.
 
         Default: ``None``.
@@ -3551,23 +3555,21 @@ def size(a, axis=None):
     >>> a = [[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
-
-    >>> a = np.asarray(a)
-    >>> np.size(a)
+    >>> np.size(a, axis=(0, 1))
     6
-    >>> np.size(a, 1)
-    3
 
     """
 
     if dpnp.is_supported_array_type(a):
         if axis is None:
             return a.size
-        return a.shape[axis]
+        _shape = a.shape
+        _axis = normalize_axis_tuple(axis, a.ndim)
+        return math.prod(_shape[ax] for ax in _axis)
 
     return numpy.size(a, axis)
 

dpnp/tests/test_manipulation.py

@@ -74,18 +74,32 @@ def test_ndim():
     assert dpnp.ndim(ia) == exp
 
 
-def test_size():
-    a = [[1, 2, 3], [4, 5, 6]]
-    ia = dpnp.array(a)
+class TestSize:
+    def test_size(self):
+        a = [[1, 2, 3], [4, 5, 6]]
+        ia = dpnp.array(a)
+
+        exp = numpy.size(a)
+        assert ia.size == exp
+        assert dpnp.size(a) == exp
+        assert dpnp.size(ia) == exp
 
-    exp = numpy.size(a)
-    assert ia.size == exp
-    assert dpnp.size(a) == exp
-    assert dpnp.size(ia) == exp
+        exp = numpy.size(a, 0)
+        assert dpnp.size(a, 0) == exp
+        assert dpnp.size(ia, 0) == exp
+
+        assert dpnp.size(ia, 1) == numpy.size(a, 1)
+
+    # TODO: include commented code in the test when numpy-2.4 is released
+    # @testing.with_requires("numpy>=2.4")
+    def test_size_tuple(self):
+        a = [[1, 2, 3], [4, 5, 6]]
+        ia = dpnp.array(a)
 
-    exp = numpy.size(a, 0)
-    assert dpnp.size(a, 0) == exp
-    assert dpnp.size(ia, 0) == exp
+        assert dpnp.size(ia, ()) == 1  # numpy.size(a, ())
+        assert dpnp.size(ia, (0,)) == 2  # numpy.size(a, (0,))
+        assert dpnp.size(ia, (1,)) == 3  # numpy.size(a, (1,))
+        assert dpnp.size(ia, (0, 1)) == 6  # numpy.size(a, (0, 1))
 
 
 class TestAppend: