add test for cp.size
@@ -2,6 +2,7 @@ import unittest
import cupy
from cupy.core import core
+from cupy import testing
class TestGetStridesForNocopyReshape(unittest.TestCase):
@@ -19,3 +20,30 @@ class TestGetStridesForNocopyReshape(unittest.TestCase):
def test_normal(self):
# TODO(nno): write test for normal case
pass
+
+
+class TestSize(unittest.TestCase):
+
+ @testing.for_all_dtypes()
+ @testing.numpy_cupy_equal()
+ def test_size(self, xp, dtype):
+ a = xp.ndarray((2, 3), dtype=dtype)
+ return xp.size(a)
+
+ @testing.for_all_dtypes()
+ @testing.numpy_cupy_equal()
+ def test_size_axis(self, xp, dtype):
+ a = xp.ndarray((2, 3), dtype=dtype)
+ return xp.size(a, axis=1)
+
+ @testing.for_all_dtypes()
+ @testing.numpy_cupy_raises()
+ def test_size_axis_error(self, xp, dtype):
+ a = xp.ndarray((2, 3), dtype=dtype)
+ return xp.size(a, axis=3)
+
+ @testing.numpy_cupy_equal()
+ @testing.slow
+ def test_size_huge(self, xp):
+ a = xp.ndarray(2 ** 32, 'b') # 4 GiB
+ return xp.size(a)
