Skip to content

Commit

Permalink
implement dpnp.ndim and dpnp.size (#2014)
Browse files Browse the repository at this point in the history
* implement dpnp.ndim

* rename direcotry

* implement dpnp.size

* remove parameters section
  • Loading branch information
vtavana authored Aug 27, 2024
1 parent 75bb4e7 commit 2f9b2fa
Show file tree
Hide file tree
Showing 5 changed files with 243 additions and 7 deletions.
1 change: 1 addition & 0 deletions .github/workflows/conda-package.yml
Original file line number Diff line number Diff line change
Expand Up @@ -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'))))"
Expand Down
52 changes: 50 additions & 2 deletions dpnp/dpnp_array.py
Original file line number Diff line number Diff line change
Expand Up @@ -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

Expand Down Expand Up @@ -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):
Expand Down
107 changes: 102 additions & 5 deletions dpnp/dpnp_iface_manipulation.py
Original file line number Diff line number Diff line change
Expand Up @@ -66,6 +66,7 @@
"flipud",
"hstack",
"moveaxis",
"ndim",
"ravel",
"repeat",
"reshape",
Expand All @@ -74,6 +75,7 @@
"rollaxis",
"row_stack",
"shape",
"size",
"squeeze",
"stack",
"swapaxes",
Expand Down Expand Up @@ -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.
Expand Down Expand Up @@ -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)
()
"""
Expand All @@ -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`.
Expand Down
24 changes: 24 additions & 0 deletions tests/test_manipulation.py
Original file line number Diff line number Diff line change
Expand Up @@ -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",
Expand Down
66 changes: 66 additions & 0 deletions tests/third_party/cupy/test_ndim.py
Original file line number Diff line number Diff line change
@@ -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

0 comments on commit 2f9b2fa

Please sign in to comment.