Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

added np.size() overload and added tests #9504

Merged
merged 11 commits into from
Apr 16, 2024
4 changes: 4 additions & 0 deletions docs/upcoming_changes/9504.np_support.rst
Original file line number Diff line number Diff line change
@@ -0,0 +1,4 @@
Added support for ``np.size()``
-------------------------------

Added ``np.size()`` support for numpy, which was previously unsupported.
10 changes: 10 additions & 0 deletions numba/np/arrayobj.py
Original file line number Diff line number Diff line change
Expand Up @@ -2555,6 +2555,16 @@ def impl(a):
return np.asarray(a).shape
return impl


@overload(np.size)
def np_size(a):
if not type_can_asarray(a):
raise errors.TypingError("The argument to np.size must be array-like")

def impl(a):
return np.asarray(a).size
return impl

# ------------------------------------------------------------------------------


Expand Down
27 changes: 27 additions & 0 deletions numba/tests/test_array_manipulation.py
Original file line number Diff line number Diff line change
Expand Up @@ -174,6 +174,9 @@ def numpy_fill_diagonal(arr, val, wrap=False):
def numpy_shape(arr):
return np.shape(arr)

def numpy_size(arr):
shourya5 marked this conversation as resolved.
Show resolved Hide resolved
return np.size(arr)


def numpy_flatnonzero(a):
return np.flatnonzero(a)
Expand Down Expand Up @@ -1219,6 +1222,30 @@ def check(x):
self.assertIn("The argument to np.shape must be array-like",
str(raises.exception))

def test_size(self):
pyfunc = numpy_size
cfunc = jit(nopython=True)(pyfunc)

def check(x):
expected = pyfunc(x)
got = cfunc(x)
self.assertPreciseEqual(got, expected)

# check arrays
for t in [(), (1,), (2, 3,), (4, 5, 6)]:
arr = np.empty(t)
check(arr)

# check scalar values
for t in [1, False, 3.14, np.int8(4), np.float32(2.718)]:
check(t)

with self.assertRaises(TypingError) as raises:
cfunc('a')

self.assertIn("The argument to np.size must be array-like",
str(raises.exception))

def test_flatnonzero_basic(self):
pyfunc = numpy_flatnonzero
cfunc = jit(nopython=True)(pyfunc)
Expand Down
18 changes: 18 additions & 0 deletions numba/tests/test_record_dtype.py
Original file line number Diff line number Diff line change
Expand Up @@ -332,6 +332,10 @@ def get_shape(rec):
return np.shape(rec.j)


def get_size(rec):
return np.size(rec.j)
shourya5 marked this conversation as resolved.
Show resolved Hide resolved


def get_charseq(ary, i):
return ary[i].n

Expand Down Expand Up @@ -1616,6 +1620,20 @@ def test_shape(self):
arr_res = cfunc(arg)
np.testing.assert_equal(arr_res, arr_expected)

def test_size(self):
# test getting the size of a nestedarray inside a record
nbarr = np.recarray(2, dtype=recordwith2darray)
nbarr[0] = np.array([(1, ((1, 2), (4, 5), (2, 3)))],
dtype=recordwith2darray)[0]

arg = nbarr[0]
pyfunc = get_size
ty = typeof(arg)
arr_expected = pyfunc(arg)
cfunc = self.get_cfunc(pyfunc, (ty,))
arr_res = cfunc(arg)
np.testing.assert_equal(arr_res, arr_expected)

def test_corner_slice(self):
# testing corner cases while slicing nested arrays
nbarr = np.recarray((1, 2, 3, 5, 7, 13, 17), dtype=recordwith4darray,
Expand Down