Merge branch 'master' into enable-test-dcptl-2121

Author: antonwolfy

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
 

conda-recipe/meta.yaml

@@ -1,6 +1,6 @@
 {% set max_compiler_and_mkl_version = environ.get("MAX_BUILD_CMPL_MKL_VERSION", "2026.0a0") %}
 {% set required_compiler_and_mkl_version = "2025.0" %}
-{% set required_dpctl_version = "0.20.0*" %}
+{% set required_dpctl_version = "0.21.0*" %}
 
 {% set pyproject = load_file_data('pyproject.toml') %}
 {% set py_build_deps = pyproject.get('build-system', {}).get('requires', []) %}

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:

environments/dpctl_pkg.txt

@@ -1,2 +1,2 @@
 --index-url https://pypi.anaconda.org/dppy/label/dev/simple
-dpctl>=0.20.0dev0
+dpctl>=0.21.0dev0

environments/dpctl_pkg.yml

@@ -2,4 +2,4 @@ name: Install dpctl package
 channels:
   - dppy/label/dev
 dependencies:
-  - dpctl>=0.20.0dev0
+  - dpctl>=0.21.0dev0

pyproject.toml

@@ -50,7 +50,7 @@ dependencies = [
   # "dpcpp-cpp-rt>=0.59.0",
   # "intel-cmplr-lib-rt>=0.59.0"
   # WARNING: use the latest dpctl dev version, otherwise stable w/f will fail
-  "dpctl>=0.20.0dev0",
+  "dpctl>=0.21.0dev0",
   "numpy>=1.25.0"
 ]
 description = "Data Parallel Extension for NumPy"