|
| 1 | +from __future__ import annotations |
| 2 | + |
1 | 3 | import itertools |
2 | 4 |
|
3 | 5 | import numpy |
@@ -435,6 +437,15 @@ def test_invalid_adv_getitem(self): |
435 | 437 | a[self.indexes] |
436 | 438 |
|
437 | 439 |
|
| 440 | +class TestArrayBadDTypeIndexAdvGetitem: |
| 441 | + @pytest.mark.parametrize("dtype", [object, "i,i", "float32", "str"]) |
| 442 | + def test_bad_dtype_adv_getitem(self, dtype): |
| 443 | + # Test various bad dtypes, supported by CuPy or not. |
| 444 | + a = cupy.arange(10) |
| 445 | + with pytest.raises(IndexError, match="arrays used as indices"): |
| 446 | + a[numpy.array([1, 2], dtype=dtype)] |
| 447 | + |
| 448 | + |
438 | 449 | @testing.parameterize( |
439 | 450 | {"shape": (0,), "indexes": ([False],)}, |
440 | 451 | { |
@@ -950,6 +961,60 @@ class TestArrayAdvancedIndexingSetitemTranspose: |
950 | 961 | def test_adv_setitem_transp(self, xp): |
951 | 962 | shape = (2, 3, 4) |
952 | 963 | a = xp.zeros(shape).transpose(0, 2, 1) |
953 | | - slices = (xp.array([1, 0]), slice(None), xp.array([2, 1])) |
| 964 | + slices = (numpy.array([1, 0]), slice(None), numpy.array([2, 1])) |
954 | 965 | a[slices] = 1 |
955 | 966 | return a |
| 967 | + |
| 968 | + |
| 969 | +class TestHugeArrays: |
| 970 | + # These tests require a lot of memory |
| 971 | + @testing.slow |
| 972 | + def test_advanced(self): |
| 973 | + try: |
| 974 | + arr = cupy.ones((1, 2**30), dtype=cupy.int8) |
| 975 | + idx = cupy.zeros(3, dtype=cupy.int32) |
| 976 | + res = arr[idx, :] |
| 977 | + # sanity check, we mostly care about it not crashing. |
| 978 | + assert res.sum() == 3 * 2**30 |
| 979 | + del res |
| 980 | + |
| 981 | + arr[idx, :] = cupy.array([[3], [3], [3]], dtype=cupy.int8) |
| 982 | + # Check 3 got written (order may not be strictly guaranteed) |
| 983 | + assert arr.sum() == 2**30 * 3 |
| 984 | + except MemoryError: |
| 985 | + pytest.skip("out of memory in test.") |
| 986 | + |
| 987 | + @testing.slow |
| 988 | + def test_take_array(self): |
| 989 | + try: |
| 990 | + arr = cupy.ones((1, 2**32), dtype=cupy.int8) |
| 991 | + arr[0, 2**30] = 0 # We should see each of these once |
| 992 | + arr[0, -1] = 0 |
| 993 | + res = arr.take(cupy.array([0, 0]), axis=0) |
| 994 | + # sanity check, we mostly care about it not crashing. |
| 995 | + assert res.sum() == 2 * (2**32 - 2) |
| 996 | + except MemoryError: |
| 997 | + pytest.skip("out of memory in test.") |
| 998 | + |
| 999 | + @testing.slow |
| 1000 | + def test_take_scalar(self): |
| 1001 | + try: |
| 1002 | + arr = cupy.ones((1, 2**32), dtype=cupy.int8) |
| 1003 | + arr[0, 2**30] = 0 # We should see each of these once |
| 1004 | + arr[0, -1] = 0 |
| 1005 | + res = arr.take(0, axis=0) |
| 1006 | + # sanity check, we mostly care about it not crashing. |
| 1007 | + assert res.sum() == 2**32 - 2 |
| 1008 | + except MemoryError: |
| 1009 | + pytest.skip("out of memory in test.") |
| 1010 | + |
| 1011 | + @testing.slow |
| 1012 | + def test_choose(self): |
| 1013 | + try: |
| 1014 | + choices = cupy.zeros((2, 2**31), dtype=cupy.int8) |
| 1015 | + choices[1, :] = 1 |
| 1016 | + res = choices[1, :].choose(choices) |
| 1017 | + # sanity check, we mostly care about it not crashing. |
| 1018 | + assert res.sum() == 2**31 |
| 1019 | + except MemoryError: |
| 1020 | + pytest.skip("out of memory in test.") |
0 commit comments