diff --git a/CHANGELOG.md b/CHANGELOG.md index 8263db169c22..2acb9cd4a5eb 100644 --- a/CHANGELOG.md +++ b/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 diff --git a/dpnp/dpnp_iface_manipulation.py b/dpnp/dpnp_iface_manipulation.py index 2806f621b5c6..725cae779eda 100644 --- a/dpnp/dpnp_iface_manipulation.py +++ b/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) diff --git a/dpnp/tests/test_manipulation.py b/dpnp/tests/test_manipulation.py index 5fc3d7df8c5f..68e8d15f24d5 100644 --- a/dpnp/tests/test_manipulation.py +++ b/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: