implement dpnp.ndim and dpnp.size

.github/workflows/conda-package.yml

@@ -62,6 +62,7 @@ env:
       third_party/cupy/sorting_tests
       third_party/cupy/statistics_tests/test_histogram.py
       third_party/cupy/statistics_tests/test_meanvar.py
+      third_party/cupy/test_ndim.py
   VER_JSON_NAME: 'version.json'
   VER_SCRIPT1: "import json; f = open('version.json', 'r'); j = json.load(f); f.close(); "
   VER_SCRIPT2: "d = j['dpnp'][0]; print('='.join((d[s] for s in ('version', 'build'))))"

dpnp/dpnp_array.py

@@ -1115,7 +1115,33 @@ def nbytes(self):
 
     @property
     def ndim(self):
-        """Number of array dimensions."""
+        """
+        Return the number of dimensions of an array.
+
+        For full documentation refer to :obj:`numpy.ndarray.ndim`.
+
+        Returns
+        -------
+        number_of_dimensions : int
+            The number of dimensions in `a`.
+
+        See Also
+        --------
+        :obj:`dpnp.ndim` : Equivalent method for any array-like input.
+        :obj:`dpnp.shape` : Return the shape of an array.
+        :obj:`dpnp.ndarray.shape` : Return the shape of an array.
+
+        Examples
+        --------
+        >>> import dpnp as np
+        >>> x = np.array([1, 2, 3])
+        >>> x.ndim
+        1
+        >>> y = np.zeros((2, 3, 4))
+        >>> y.ndim
+        3
+
+        """
 
         return self._array_obj.ndim
 
@@ -1389,7 +1415,29 @@ def shape(self, newshape):
 
     @property
     def size(self):
-        """Number of elements in the array."""
+        """
+        Number of elements in the array.
+
+        Returns
+        -------
+        element_count : int
+            Number of elements in the array.
+
+        See Also
+        --------
+        :obj:`dpnp.size` : Return the number of elements along a given axis.
+        :obj:`dpnp.shape` : Return the shape of an array.
+        :obj:`dpnp.ndarray.shape` : Return the shape of an array.
+
+        Examples
+        --------
+        >>> import dpnp as np
+        >>> x = np.zeros((3, 5, 2), dtype=np.complex64)
+        >>> x.size
+        30
+
+        """
+
         return self._array_obj.size
 
     def sort(self, axis=-1, kind=None, order=None):

dpnp/dpnp_iface_manipulation.py

@@ -66,6 +66,7 @@
     "flipud",
     "hstack",
     "moveaxis",
+    "ndim",
     "ravel",
     "repeat",
     "reshape",
@@ -74,6 +75,7 @@
     "rollaxis",
     "row_stack",
     "shape",
+    "size",
     "squeeze",
     "stack",
     "swapaxes",
@@ -1357,6 +1359,48 @@ def moveaxis(a, source, destination):
     )
 
 
+def ndim(a):
+    """
+    Return the number of dimensions of array-like input.
+
+    For full documentation refer to :obj:`numpy.ndim`.
+
+    Parameters
+    ----------
+    a : array_like
+        Input data.
+
+    Returns
+    -------
+    number_of_dimensions : int
+        The number of dimensions in `a`. Scalars are zero-dimensional.
+
+    See Also
+    --------
+    :obj:`dpnp.ndarray.ndim` : Equivalent method for `dpnp.ndarray`
+                        or `usm_ndarray` input.
+    :obj:`dpnp.shape` : Return the shape of an array.
+    :obj:`dpnp.ndarray.shape` : Return the shape of an array.
+
+    Examples
+    --------
+    >>> import dpnp as np
+    >>> a = [[1, 2, 3], [4, 5, 6]]
+    >>> np.ndim(a)
+    2
+    >>> a = np.asarray(a)
+    >>> np.ndim(a)
+    2
+    >>> np.ndim(1)
+    0
+
+    """
+
+    if dpnp.is_supported_array_type(a):
+        return a.ndim
+    return numpy.ndim(a)
+
+
 def ravel(a, order="C"):
     """
     Return a contiguous flattened array.
@@ -1739,14 +1783,14 @@ def shape(a):
 
     Examples
     --------
-    >>> import dpnp as dp
-    >>> dp.shape(dp.eye(3))
+    >>> import dpnp as np
+    >>> np.shape(np.eye(3))
     (3, 3)
-    >>> dp.shape([[1, 3]])
+    >>> np.shape([[1, 3]])
     (1, 2)
-    >>> dp.shape([0])
+    >>> np.shape([0])
     (1,)
-    >>> dp.shape(0)
+    >>> np.shape(0)
     ()
 
     """
@@ -1756,6 +1800,59 @@ def shape(a):
     return numpy.shape(a)
 
 
+def size(a, axis=None):
+    """
+    Return the number of elements along a given axis.
+
+    For full documentation refer to :obj:`numpy.size`.
+
+    Parameters
+    ----------
+    a : array_like
+        Input data.
+    axis : {None, int}, optional
+        Axis along which the elements are counted.
+        By default, give the total number of elements.
+        Default: ``None``.
+
+    Returns
+    -------
+    element_count : int
+        Number of elements along the specified axis.
+
+    See Also
+    --------
+    :obj:`dpnp.ndarray.size` : number of elements in array.
+    :obj:`dpnp.shape` : Return the shape of an array.
+    :obj:`dpnp.ndarray.shape` : Return the shape of an array.
+
+    Examples
+    --------
+    >>> import dpnp as np
+    >>> a = [[1, 2, 3], [4, 5, 6]]
+    >>> np.size(a)
+    6
+    >>> np.size(a, 1)
+    3
+    >>> np.size(a, 0)
+    2
+
+    >>> a = np.asarray(a)
+    >>> np.size(a)
+    6
+    >>> np.size(a, 1)
+    3
+
+    """
+
+    if dpnp.is_supported_array_type(a):
+        if axis is None:
+            return a.size
+        return a.shape[axis]
+
+    return numpy.size(a, axis)
+
+
 def squeeze(a, /, axis=None):
     """
     Removes singleton dimensions (axes) from array `a`.

tests/test_manipulation.py

@@ -89,6 +89,30 @@ def test_result_type_only_arrays():
     assert dpnp.result_type(*X) == numpy.result_type(*X_np)
 
 
+def test_ndim():
+    a = [[1, 2, 3], [4, 5, 6]]
+    ia = dpnp.array(a)
+
+    exp = numpy.ndim(a)
+    assert ia.ndim == exp
+    assert dpnp.ndim(a) == exp
+    assert dpnp.ndim(ia) == exp
+
+
+def test_size():
+    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, 0)
+    assert dpnp.size(a, 0) == exp
+    assert dpnp.size(ia, 0) == exp
+
+
 class TestRepeat:
     @pytest.mark.parametrize(
         "data",

tests/third_party/cupy/test_ndim.py

@@ -0,0 +1,66 @@
+import unittest
+
+import numpy
+
+import dpnp as cupy
+from tests.third_party.cupy import testing
+
+
+class TestNdim(unittest.TestCase):
+    @testing.numpy_cupy_equal()
+    def test_ndim_ndarray1d(self, xp):
+        return xp.ndim(xp.arange(5))
+
+    @testing.numpy_cupy_equal()
+    def test_ndim_ndarray2d(self, xp):
+        return xp.ndim(xp.ones((2, 4)))
+
+    @testing.numpy_cupy_equal()
+    def test_ndim_ndarray0d(self, xp):
+        return xp.ndim(xp.asarray(5))
+
+    @testing.numpy_cupy_equal()
+    def test_ndim_scalar(self, xp):
+        return xp.ndim(5)
+
+    @testing.numpy_cupy_equal()
+    def test_ndim_none(self, xp):
+        return xp.ndim(None)
+
+    @testing.numpy_cupy_equal()
+    def test_ndim_string(self, xp):
+        return xp.ndim("abc")
+
+    @testing.numpy_cupy_equal()
+    def test_ndim_list1(self, xp):
+        return xp.ndim([1, 2, 3])
+
+    @testing.numpy_cupy_equal()
+    def test_ndim_list2(self, xp):
+        return xp.ndim([[1, 2, 3], [4, 5, 6]])
+
+    @testing.numpy_cupy_equal()
+    def test_ndim_tuple(self, xp):
+        return xp.ndim(((1, 2, 3), (4, 5, 6)))
+
+    @testing.numpy_cupy_equal()
+    def test_ndim_set(self, xp):
+        return xp.ndim({1, 2, 3})
+
+    @testing.numpy_cupy_equal()
+    def test_ndim_object(self, xp):
+        return xp.ndim(dict(a=5, b="b"))
+
+    # numpy.dim works on dpnp arrays and dpnp.ndim works on NumPy arrays
+    def test_ndim_array_function(self):
+        a = cupy.ones((4, 4))
+        assert numpy.ndim(a) == 2
+
+        a = cupy.asarray(5)
+        assert numpy.ndim(a) == 0
+
+        a = numpy.ones((4, 4))
+        assert cupy.ndim(a) == 2
+
+        a = numpy.asarray(5)
+        assert cupy.ndim(a) == 0