• Home
  • Line#
  • Scopes#
  • Navigate#
  • Raw
  • Download
1# Owner(s): ["module: dynamo"]
2
3import functools
4import operator
5import re
6import sys
7import warnings
8from itertools import product
9from unittest import expectedFailure as xfail, skipIf as skipif, SkipTest
10
11import pytest
12from pytest import raises as assert_raises
13
14from torch.testing._internal.common_utils import (
15    instantiate_parametrized_tests,
16    parametrize,
17    run_tests,
18    skipIfTorchDynamo,
19    TEST_WITH_TORCHDYNAMO,
20    TestCase,
21    xpassIfTorchDynamo,
22)
23
24
25if TEST_WITH_TORCHDYNAMO:
26    import numpy as np
27    from numpy.testing import (
28        assert_,
29        assert_array_equal,
30        assert_equal,
31        assert_warns,
32        HAS_REFCOUNT,
33    )
34else:
35    import torch._numpy as np
36    from torch._numpy.testing import (
37        assert_,
38        assert_array_equal,
39        assert_equal,
40        assert_warns,
41        HAS_REFCOUNT,
42    )
43
44
45skip = functools.partial(skipif, True)
46
47
48@instantiate_parametrized_tests
49class TestIndexing(TestCase):
50    def test_index_no_floats(self):
51        a = np.array([[[5]]])
52
53        assert_raises(IndexError, lambda: a[0.0])
54        assert_raises(IndexError, lambda: a[0, 0.0])
55        assert_raises(IndexError, lambda: a[0.0, 0])
56        assert_raises(IndexError, lambda: a[0.0, :])
57        assert_raises(IndexError, lambda: a[:, 0.0])
58        assert_raises(IndexError, lambda: a[:, 0.0, :])
59        assert_raises(IndexError, lambda: a[0.0, :, :])
60        assert_raises(IndexError, lambda: a[0, 0, 0.0])
61        assert_raises(IndexError, lambda: a[0.0, 0, 0])
62        assert_raises(IndexError, lambda: a[0, 0.0, 0])
63        assert_raises(IndexError, lambda: a[-1.4])
64        assert_raises(IndexError, lambda: a[0, -1.4])
65        assert_raises(IndexError, lambda: a[-1.4, 0])
66        assert_raises(IndexError, lambda: a[-1.4, :])
67        assert_raises(IndexError, lambda: a[:, -1.4])
68        assert_raises(IndexError, lambda: a[:, -1.4, :])
69        assert_raises(IndexError, lambda: a[-1.4, :, :])
70        assert_raises(IndexError, lambda: a[0, 0, -1.4])
71        assert_raises(IndexError, lambda: a[-1.4, 0, 0])
72        assert_raises(IndexError, lambda: a[0, -1.4, 0])
73        # Note torch validates index arguments "depth-first", so will prioritise
74        # raising TypeError over IndexError, e.g.
75        #
76        #     >>> a = np.array([[[5]]])
77        #     >>> a[0.0:, 0.0]
78        #     IndexError: only integers, slices (`:`), ellipsis (`...`),
79        #     numpy.newaxis #     (`None`) and integer or boolean arrays are
80        #     valid indices
81        #     >>> t = torch.as_tensor([[[5]]])  # identical to a
82        #     >>> t[0.0:, 0.0]
83        #     TypeError: slice indices must be integers or None or have an
84        #     __index__ method
85        #
86        assert_raises((IndexError, TypeError), lambda: a[0.0:, 0.0])
87        assert_raises((IndexError, TypeError), lambda: a[0.0:, 0.0, :])
88
89    def test_slicing_no_floats(self):
90        a = np.array([[5]])
91
92        # start as float.
93        assert_raises(TypeError, lambda: a[0.0:])
94        assert_raises(TypeError, lambda: a[0:, 0.0:2])
95        assert_raises(TypeError, lambda: a[0.0::2, :0])
96        assert_raises(TypeError, lambda: a[0.0:1:2, :])
97        assert_raises(TypeError, lambda: a[:, 0.0:])
98        # stop as float.
99        assert_raises(TypeError, lambda: a[:0.0])
100        assert_raises(TypeError, lambda: a[:0, 1:2.0])
101        assert_raises(TypeError, lambda: a[:0.0:2, :0])
102        assert_raises(TypeError, lambda: a[:0.0, :])
103        assert_raises(TypeError, lambda: a[:, 0:4.0:2])
104        # step as float.
105        assert_raises(TypeError, lambda: a[::1.0])
106        assert_raises(TypeError, lambda: a[0:, :2:2.0])
107        assert_raises(TypeError, lambda: a[1::4.0, :0])
108        assert_raises(TypeError, lambda: a[::5.0, :])
109        assert_raises(TypeError, lambda: a[:, 0:4:2.0])
110        # mixed.
111        assert_raises(TypeError, lambda: a[1.0:2:2.0])
112        assert_raises(TypeError, lambda: a[1.0::2.0])
113        assert_raises(TypeError, lambda: a[0:, :2.0:2.0])
114        assert_raises(TypeError, lambda: a[1.0:1:4.0, :0])
115        assert_raises(TypeError, lambda: a[1.0:5.0:5.0, :])
116        assert_raises(TypeError, lambda: a[:, 0.4:4.0:2.0])
117        # should still get the DeprecationWarning if step = 0.
118        assert_raises(TypeError, lambda: a[::0.0])
119
120    @skip(reason="torch allows slicing with non-0d array components")
121    def test_index_no_array_to_index(self):
122        # No non-scalar arrays.
123        a = np.array([[[1]]])
124
125        assert_raises(TypeError, lambda: a[a:a:a])
126        # Conversely, using scalars doesn't raise in NumPy, e.g.
127        #
128        #     >>> i = np.int64(1)
129        #     >>> a[i:i:i]
130        #     array([], shape=(0, 1, 1), dtype=int64)
131        #
132
133    def test_none_index(self):
134        # `None` index adds newaxis
135        a = np.array([1, 2, 3])
136        assert_equal(a[None], a[np.newaxis])
137        assert_equal(a[None].ndim, a.ndim + 1)
138
139    @skip
140    def test_empty_tuple_index(self):
141        # Empty tuple index creates a view
142        a = np.array([1, 2, 3])
143        assert_equal(a[()], a)
144        assert_(a[()].tensor._base is a.tensor)
145        a = np.array(0)
146        assert_(isinstance(a[()], np.int_))
147
148    def test_same_kind_index_casting(self):
149        # Indexes should be cast with same-kind and not safe, even if that
150        # is somewhat unsafe. So test various different code paths.
151        index = np.arange(5)
152        u_index = index.astype(np.uint8)  # i.e. cast to default uint indexing dtype
153        arr = np.arange(10)
154
155        assert_array_equal(arr[index], arr[u_index])
156        arr[u_index] = np.arange(5)
157        assert_array_equal(arr, np.arange(10))
158
159        arr = np.arange(10).reshape(5, 2)
160        assert_array_equal(arr[index], arr[u_index])
161
162        arr[u_index] = np.arange(5)[:, None]
163        assert_array_equal(arr, np.arange(5)[:, None].repeat(2, axis=1))
164
165        arr = np.arange(25).reshape(5, 5)
166        assert_array_equal(arr[u_index, u_index], arr[index, index])
167
168    def test_empty_fancy_index(self):
169        # Empty list index creates an empty array
170        # with the same dtype (but with weird shape)
171        a = np.array([1, 2, 3])
172        assert_equal(a[[]], [])
173        assert_equal(a[[]].dtype, a.dtype)
174
175        b = np.array([], dtype=np.intp)
176        assert_equal(a[[]], [])
177        assert_equal(a[[]].dtype, a.dtype)
178
179        b = np.array([])
180        assert_raises(IndexError, a.__getitem__, b)
181
182    def test_ellipsis_index(self):
183        a = np.array([[1, 2, 3], [4, 5, 6], [7, 8, 9]])
184        assert_(a[...] is not a)
185        assert_equal(a[...], a)
186        # `a[...]` was `a` in numpy <1.9.
187
188        # Slicing with ellipsis can skip an
189        # arbitrary number of dimensions
190        assert_equal(a[0, ...], a[0])
191        assert_equal(a[0, ...], a[0, :])
192        assert_equal(a[..., 0], a[:, 0])
193
194        # Slicing with ellipsis always results
195        # in an array, not a scalar
196        assert_equal(a[0, ..., 1], np.array(2))
197
198        # Assignment with `(Ellipsis,)` on 0-d arrays
199        b = np.array(1)
200        b[(Ellipsis,)] = 2
201        assert_equal(b, 2)
202
203    @xpassIfTorchDynamo  # 'torch_.np.array() does not have base attribute.
204    def test_ellipsis_index_2(self):
205        a = np.array([[1, 2, 3], [4, 5, 6], [7, 8, 9]])
206        assert_(a[...] is not a)
207        assert_equal(a[...], a)
208        # `a[...]` was `a` in numpy <1.9.
209        assert_(a[...].base is a)
210
211    def test_single_int_index(self):
212        # Single integer index selects one row
213        a = np.array([[1, 2, 3], [4, 5, 6], [7, 8, 9]])
214
215        assert_equal(a[0], [1, 2, 3])
216        assert_equal(a[-1], [7, 8, 9])
217
218        # Index out of bounds produces IndexError
219        assert_raises(IndexError, a.__getitem__, 1 << 30)
220        # Index overflow produces IndexError
221        # Note torch raises RuntimeError here
222        assert_raises((IndexError, RuntimeError), a.__getitem__, 1 << 64)
223
224    def test_single_bool_index(self):
225        # Single boolean index
226        a = np.array([[1, 2, 3], [4, 5, 6], [7, 8, 9]])
227
228        assert_equal(a[np.array(True)], a[None])
229        assert_equal(a[np.array(False)], a[None][0:0])
230
231    def test_boolean_shape_mismatch(self):
232        arr = np.ones((5, 4, 3))
233
234        index = np.array([True])
235        assert_raises(IndexError, arr.__getitem__, index)
236
237        index = np.array([False] * 6)
238        assert_raises(IndexError, arr.__getitem__, index)
239
240        index = np.zeros((4, 4), dtype=bool)
241        assert_raises(IndexError, arr.__getitem__, index)
242
243        assert_raises(IndexError, arr.__getitem__, (slice(None), index))
244
245    def test_boolean_indexing_onedim(self):
246        # Indexing a 2-dimensional array with
247        # boolean array of length one
248        a = np.array([[0.0, 0.0, 0.0]])
249        b = np.array([True], dtype=bool)
250        assert_equal(a[b], a)
251        # boolean assignment
252        a[b] = 1.0
253        assert_equal(a, [[1.0, 1.0, 1.0]])
254
255    @skip(reason="NP_VER: fails on CI")
256    def test_boolean_assignment_value_mismatch(self):
257        # A boolean assignment should fail when the shape of the values
258        # cannot be broadcast to the subscription. (see also gh-3458)
259        a = np.arange(4)
260
261        def f(a, v):
262            a[a > -1] = v
263
264        assert_raises((RuntimeError, ValueError, TypeError), f, a, [])
265        assert_raises((RuntimeError, ValueError, TypeError), f, a, [1, 2, 3])
266        assert_raises((RuntimeError, ValueError, TypeError), f, a[:1], [1, 2, 3])
267
268    def test_boolean_indexing_twodim(self):
269        # Indexing a 2-dimensional array with
270        # 2-dimensional boolean array
271        a = np.array([[1, 2, 3], [4, 5, 6], [7, 8, 9]])
272        b = np.array([[True, False, True], [False, True, False], [True, False, True]])
273        assert_equal(a[b], [1, 3, 5, 7, 9])
274        assert_equal(a[b[1]], [[4, 5, 6]])
275        assert_equal(a[b[0]], a[b[2]])
276
277        # boolean assignment
278        a[b] = 0
279        assert_equal(a, [[0, 2, 0], [4, 0, 6], [0, 8, 0]])
280
281    def test_boolean_indexing_list(self):
282        # Regression test for #13715. It's a use-after-free bug which the
283        # test won't directly catch, but it will show up in valgrind.
284        a = np.array([1, 2, 3])
285        b = [True, False, True]
286        # Two variants of the test because the first takes a fast path
287        assert_equal(a[b], [1, 3])
288        assert_equal(a[None, b], [[1, 3]])
289
290    def test_reverse_strides_and_subspace_bufferinit(self):
291        # This tests that the strides are not reversed for simple and
292        # subspace fancy indexing.
293        a = np.ones(5)
294        b = np.zeros(5, dtype=np.intp)[::-1]
295        c = np.arange(5)[::-1]
296
297        a[b] = c
298        # If the strides are not reversed, the 0 in the arange comes last.
299        assert_equal(a[0], 0)
300
301        # This also tests that the subspace buffer is initialized:
302        a = np.ones((5, 2))
303        c = np.arange(10).reshape(5, 2)[::-1]
304        a[b, :] = c
305        assert_equal(a[0], [0, 1])
306
307    def test_reversed_strides_result_allocation(self):
308        # Test a bug when calculating the output strides for a result array
309        # when the subspace size was 1 (and test other cases as well)
310        a = np.arange(10)[:, None]
311        i = np.arange(10)[::-1]
312        assert_array_equal(a[i], a[i.copy("C")])
313
314        a = np.arange(20).reshape(-1, 2)
315
316    def test_uncontiguous_subspace_assignment(self):
317        # During development there was a bug activating a skip logic
318        # based on ndim instead of size.
319        a = np.full((3, 4, 2), -1)
320        b = np.full((3, 4, 2), -1)
321
322        a[[0, 1]] = np.arange(2 * 4 * 2).reshape(2, 4, 2).T
323        b[[0, 1]] = np.arange(2 * 4 * 2).reshape(2, 4, 2).T.copy()
324
325        assert_equal(a, b)
326
327    @skip(reason="torch does not limit dims to 32")
328    def test_too_many_fancy_indices_special_case(self):
329        # Just documents behaviour, this is a small limitation.
330        a = np.ones((1,) * 32)  # 32 is NPY_MAXDIMS
331        assert_raises(IndexError, a.__getitem__, (np.array([0]),) * 32)
332
333    def test_scalar_array_bool(self):
334        # NumPy bools can be used as boolean index (python ones as of yet not)
335        a = np.array(1)
336        assert_equal(a[np.bool_(True)], a[np.array(True)])
337        assert_equal(a[np.bool_(False)], a[np.array(False)])
338
339        # After deprecating bools as integers:
340        # a = np.array([0,1,2])
341        # assert_equal(a[True, :], a[None, :])
342        # assert_equal(a[:, True], a[:, None])
343        #
344        # assert_(not np.may_share_memory(a, a[True, :]))
345
346    def test_everything_returns_views(self):
347        # Before `...` would return a itself.
348        a = np.arange(5)
349
350        assert_(a is not a[()])
351        assert_(a is not a[...])
352        assert_(a is not a[:])
353
354    def test_broaderrors_indexing(self):
355        a = np.zeros((5, 5))
356        assert_raises(IndexError, a.__getitem__, ([0, 1], [0, 1, 2]))
357        assert_raises(IndexError, a.__setitem__, ([0, 1], [0, 1, 2]), 0)
358
359    def test_trivial_fancy_out_of_bounds(self):
360        a = np.zeros(5)
361        ind = np.ones(20, dtype=np.intp)
362        ind[-1] = 10
363        assert_raises(IndexError, a.__getitem__, ind)
364        assert_raises((IndexError, RuntimeError), a.__setitem__, ind, 0)
365        ind = np.ones(20, dtype=np.intp)
366        ind[0] = 11
367        assert_raises(IndexError, a.__getitem__, ind)
368        assert_raises((IndexError, RuntimeError), a.__setitem__, ind, 0)
369
370    def test_trivial_fancy_not_possible(self):
371        # Test that the fast path for trivial assignment is not incorrectly
372        # used when the index is not contiguous or 1D, see also gh-11467.
373        a = np.arange(6)
374        idx = np.arange(6, dtype=np.intp).reshape(2, 1, 3)[:, :, 0]
375        assert_array_equal(a[idx], idx)
376
377        # this case must not go into the fast path, note that idx is
378        # a non-contiuguous none 1D array here.
379        a[idx] = -1
380        res = np.arange(6)
381        res[0] = -1
382        res[3] = -1
383        assert_array_equal(a, res)
384
385    def test_memory_order(self):
386        # This is not necessary to preserve. Memory layouts for
387        # more complex indices are not as simple.
388        a = np.arange(10)
389        b = np.arange(10).reshape(5, 2).T
390        assert_(a[b].flags.f_contiguous)
391
392        # Takes a different implementation branch:
393        a = a.reshape(-1, 1)
394        assert_(a[b, 0].flags.f_contiguous)
395
396    @skipIfTorchDynamo()  # XXX: flaky, depends on implementation details
397    def test_small_regressions(self):
398        # Reference count of intp for index checks
399        a = np.array([0])
400        if HAS_REFCOUNT:
401            refcount = sys.getrefcount(np.dtype(np.intp))
402        # item setting always checks indices in separate function:
403        a[np.array([0], dtype=np.intp)] = 1
404        a[np.array([0], dtype=np.uint8)] = 1
405        assert_raises(IndexError, a.__setitem__, np.array([1], dtype=np.intp), 1)
406        assert_raises(IndexError, a.__setitem__, np.array([1], dtype=np.uint8), 1)
407
408        if HAS_REFCOUNT:
409            assert_equal(sys.getrefcount(np.dtype(np.intp)), refcount)
410
411    def test_tuple_subclass(self):
412        arr = np.ones((5, 5))
413
414        # A tuple subclass should also be an nd-index
415        class TupleSubclass(tuple):
416            pass
417
418        index = ([1], [1])
419        index = TupleSubclass(index)
420        assert_(arr[index].shape == (1,))
421        # Unlike the non nd-index:
422        assert_(arr[index,].shape != (1,))
423
424    @xpassIfTorchDynamo  # (reason="XXX: low-prio behaviour to support")
425    def test_broken_sequence_not_nd_index(self):
426        # See https://github.com/numpy/numpy/issues/5063
427        # If we have an object which claims to be a sequence, but fails
428        # on item getting, this should not be converted to an nd-index (tuple)
429        # If this object happens to be a valid index otherwise, it should work
430        # This object here is very dubious and probably bad though:
431        class SequenceLike:
432            def __index__(self):
433                return 0
434
435            def __len__(self):
436                return 1
437
438            def __getitem__(self, item):
439                raise IndexError("Not possible")
440
441        arr = np.arange(10)
442        assert_array_equal(arr[SequenceLike()], arr[SequenceLike(),])
443
444        # also test that field indexing does not segfault
445        # for a similar reason, by indexing a structured array
446        arr = np.zeros((1,), dtype=[("f1", "i8"), ("f2", "i8")])
447        assert_array_equal(arr[SequenceLike()], arr[SequenceLike(),])
448
449    def test_indexing_array_weird_strides(self):
450        # See also gh-6221
451        # the shapes used here come from the issue and create the correct
452        # size for the iterator buffering size.
453        x = np.ones(10)
454        x2 = np.ones((10, 2))
455        ind = np.arange(10)[:, None, None, None]
456        ind = np.broadcast_to(ind, (10, 55, 4, 4))
457
458        # single advanced index case
459        assert_array_equal(x[ind], x[ind.copy()])
460        # higher dimensional advanced index
461        zind = np.zeros(4, dtype=np.intp)
462        assert_array_equal(x2[ind, zind], x2[ind.copy(), zind])
463
464    def test_indexing_array_negative_strides(self):
465        # From gh-8264,
466        # core dumps if negative strides are used in iteration
467        arro = np.zeros((4, 4))
468        arr = arro[::-1, ::-1]
469
470        slices = (slice(None), [0, 1, 2, 3])
471        arr[slices] = 10
472        assert_array_equal(arr, 10.0)
473
474    @parametrize("index", [True, False, np.array([0])])
475    @parametrize("num", [32, 40])
476    @parametrize("original_ndim", [1, 32])
477    def test_too_many_advanced_indices(self, index, num, original_ndim):
478        # These are limitations based on the number of arguments we can process.
479        # For `num=32` (and all boolean cases), the result is actually define;
480        # but the use of NpyIter (NPY_MAXARGS) limits it for technical reasons.
481        if not (isinstance(index, np.ndarray) and original_ndim < num):
482            # unskipped cases fail because of assigning too many indices
483            raise SkipTest("torch does not limit dims to 32")
484
485        arr = np.ones((1,) * original_ndim)
486        with pytest.raises(IndexError):
487            arr[(index,) * num]
488        with pytest.raises(IndexError):
489            arr[(index,) * num] = 1.0
490
491    def test_nontuple_ndindex(self):
492        a = np.arange(25).reshape((5, 5))
493        assert_equal(a[[0, 1]], np.array([a[0], a[1]]))
494        assert_equal(a[[0, 1], [0, 1]], np.array([0, 6]))
495        raise SkipTest(
496            "torch happily consumes non-tuple sequences with multi-axis "
497            "indices (i.e. slices) as an index, whereas NumPy invalidates "
498            "them, assumedly to keep things simple. This invalidation "
499            "behaviour is just too niche to bother emulating."
500        )
501        assert_raises(IndexError, a.__getitem__, [slice(None)])
502
503
504@instantiate_parametrized_tests
505class TestBroadcastedAssignments(TestCase):
506    def assign(self, a, ind, val):
507        a[ind] = val
508        return a
509
510    def test_prepending_ones(self):
511        a = np.zeros((3, 2))
512
513        a[...] = np.ones((1, 3, 2))
514        # Fancy with subspace with and without transpose
515        a[[0, 1, 2], :] = np.ones((1, 3, 2))
516        a[:, [0, 1]] = np.ones((1, 3, 2))
517        # Fancy without subspace (with broadcasting)
518        a[[[0], [1], [2]], [0, 1]] = np.ones((1, 3, 2))
519
520    def test_prepend_not_one(self):
521        assign = self.assign
522        s_ = np.s_
523        a = np.zeros(5)
524
525        # Too large and not only ones.
526        try:
527            assign(a, s_[...], np.ones((2, 1)))
528        except Exception as e:
529            self.assertTrue(isinstance(e, (ValueError, RuntimeError)))
530        assert_raises(
531            (ValueError, RuntimeError), assign, a, s_[[1, 2, 3],], np.ones((2, 1))
532        )
533        assert_raises(
534            (ValueError, RuntimeError), assign, a, s_[[[1], [2]],], np.ones((2, 2, 1))
535        )
536
537    def test_simple_broadcasting_errors(self):
538        assign = self.assign
539        s_ = np.s_
540        a = np.zeros((5, 1))
541
542        try:
543            assign(a, s_[...], np.zeros((5, 2)))
544        except Exception as e:
545            self.assertTrue(isinstance(e, (ValueError, RuntimeError)))
546        try:
547            assign(a, s_[...], np.zeros((5, 0)))
548        except Exception as e:
549            self.assertTrue(isinstance(e, (ValueError, RuntimeError)))
550        assert_raises(
551            (ValueError, RuntimeError), assign, a, s_[:, [0]], np.zeros((5, 2))
552        )
553        assert_raises(
554            (ValueError, RuntimeError), assign, a, s_[:, [0]], np.zeros((5, 0))
555        )
556        assert_raises(
557            (ValueError, RuntimeError), assign, a, s_[[0], :], np.zeros((2, 1))
558        )
559
560    @parametrize(
561        "index", [(..., [1, 2], slice(None)), ([0, 1], ..., 0), (..., [1, 2], [1, 2])]
562    )
563    def test_broadcast_error_reports_correct_shape(self, index):
564        values = np.zeros((100, 100))  # will never broadcast below
565
566        arr = np.zeros((3, 4, 5, 6, 7))
567
568        with pytest.raises((ValueError, RuntimeError)) as e:
569            arr[index] = values
570
571        shape = arr[index].shape
572        r_inner_shape = "".join(f"{side}, ?" for side in shape[:-1]) + str(shape[-1])
573        assert re.search(rf"[\(\[]{r_inner_shape}[\]\)]$", str(e.value))
574
575    def test_index_is_larger(self):
576        # Simple case of fancy index broadcasting of the index.
577        a = np.zeros((5, 5))
578        a[[[0], [1], [2]], [0, 1, 2]] = [2, 3, 4]
579
580        assert_((a[:3, :3] == [2, 3, 4]).all())
581
582    def test_broadcast_subspace(self):
583        a = np.zeros((100, 100))
584        v = np.arange(100)[:, None]
585        b = np.arange(100)[::-1]
586        a[b] = v
587        assert_((a[::-1] == v).all())
588
589
590class TestFancyIndexingCast(TestCase):
591    @xpassIfTorchDynamo  # (
592    #    reason="XXX: low-prio to support assigning complex values on floating arrays"
593    # )
594    def test_boolean_index_cast_assign(self):
595        # Setup the boolean index and float arrays.
596        shape = (8, 63)
597        bool_index = np.zeros(shape).astype(bool)
598        bool_index[0, 1] = True
599        zero_array = np.zeros(shape)
600
601        # Assigning float is fine.
602        zero_array[bool_index] = np.array([1])
603        assert_equal(zero_array[0, 1], 1)
604
605        # Fancy indexing works, although we get a cast warning.
606        assert_warns(
607            np.ComplexWarning, zero_array.__setitem__, ([0], [1]), np.array([2 + 1j])
608        )
609        assert_equal(zero_array[0, 1], 2)  # No complex part
610
611        # Cast complex to float, throwing away the imaginary portion.
612        assert_warns(
613            np.ComplexWarning, zero_array.__setitem__, bool_index, np.array([1j])
614        )
615        assert_equal(zero_array[0, 1], 0)
616
617
618@xfail  # (reason="XXX: requires broadcast() and broadcast_to()")
619class TestMultiIndexingAutomated(TestCase):
620    """
621    These tests use code to mimic the C-Code indexing for selection.
622
623    NOTE:
624
625        * This still lacks tests for complex item setting.
626        * If you change behavior of indexing, you might want to modify
627          these tests to try more combinations.
628        * Behavior was written to match numpy version 1.8. (though a
629          first version matched 1.7.)
630        * Only tuple indices are supported by the mimicking code.
631          (and tested as of writing this)
632        * Error types should match most of the time as long as there
633          is only one error. For multiple errors, what gets raised
634          will usually not be the same one. They are *not* tested.
635
636    Update 2016-11-30: It is probably not worth maintaining this test
637    indefinitely and it can be dropped if maintenance becomes a burden.
638
639    """
640
641    def setupUp(self):
642        self.a = np.arange(np.prod([3, 1, 5, 6])).reshape(3, 1, 5, 6)
643        self.b = np.empty((3, 0, 5, 6))
644        self.complex_indices = [
645            "skip",
646            Ellipsis,
647            0,
648            # Boolean indices, up to 3-d for some special cases of eating up
649            # dimensions, also need to test all False
650            np.array([True, False, False]),
651            np.array([[True, False], [False, True]]),
652            np.array([[[False, False], [False, False]]]),
653            # Some slices:
654            slice(-5, 5, 2),
655            slice(1, 1, 100),
656            slice(4, -1, -2),
657            slice(None, None, -3),
658            # Some Fancy indexes:
659            np.empty((0, 1, 1), dtype=np.intp),  # empty and can be broadcast
660            np.array([0, 1, -2]),
661            np.array([[2], [0], [1]]),
662            np.array([[0, -1], [0, 1]], dtype=np.dtype("intp")),
663            np.array([2, -1], dtype=np.int8),
664            np.zeros([1] * 31, dtype=int),  # trigger too large array.
665            np.array([0.0, 1.0]),
666        ]  # invalid datatype
667        # Some simpler indices that still cover a bit more
668        self.simple_indices = [Ellipsis, None, -1, [1], np.array([True]), "skip"]
669        # Very simple ones to fill the rest:
670        self.fill_indices = [slice(None, None), 0]
671
672    def _get_multi_index(self, arr, indices):
673        """Mimic multi dimensional indexing.
674
675        Parameters
676        ----------
677        arr : ndarray
678            Array to be indexed.
679        indices : tuple of index objects
680
681        Returns
682        -------
683        out : ndarray
684            An array equivalent to the indexing operation (but always a copy).
685            `arr[indices]` should be identical.
686        no_copy : bool
687            Whether the indexing operation requires a copy. If this is `True`,
688            `np.may_share_memory(arr, arr[indices])` should be `True` (with
689            some exceptions for scalars and possibly 0-d arrays).
690
691        Notes
692        -----
693        While the function may mostly match the errors of normal indexing this
694        is generally not the case.
695        """
696        in_indices = list(indices)
697        indices = []
698        # if False, this is a fancy or boolean index
699        no_copy = True
700        # number of fancy/scalar indexes that are not consecutive
701        num_fancy = 0
702        # number of dimensions indexed by a "fancy" index
703        fancy_dim = 0
704        # NOTE: This is a funny twist (and probably OK to change).
705        # The boolean array has illegal indexes, but this is
706        # allowed if the broadcast fancy-indices are 0-sized.
707        # This variable is to catch that case.
708        error_unless_broadcast_to_empty = False
709
710        # We need to handle Ellipsis and make arrays from indices, also
711        # check if this is fancy indexing (set no_copy).
712        ndim = 0
713        ellipsis_pos = None  # define here mostly to replace all but first.
714        for i, indx in enumerate(in_indices):
715            if indx is None:
716                continue
717            if isinstance(indx, np.ndarray) and indx.dtype == bool:
718                no_copy = False
719                if indx.ndim == 0:
720                    raise IndexError
721                # boolean indices can have higher dimensions
722                ndim += indx.ndim
723                fancy_dim += indx.ndim
724                continue
725            if indx is Ellipsis:
726                if ellipsis_pos is None:
727                    ellipsis_pos = i
728                    continue  # do not increment ndim counter
729                raise IndexError
730            if isinstance(indx, slice):
731                ndim += 1
732                continue
733            if not isinstance(indx, np.ndarray):
734                # This could be open for changes in numpy.
735                # numpy should maybe raise an error if casting to intp
736                # is not safe. It rejects np.array([1., 2.]) but not
737                # [1., 2.] as index (same for ie. np.take).
738                # (Note the importance of empty lists if changing this here)
739                try:
740                    indx = np.array(indx, dtype=np.intp)
741                except ValueError:
742                    raise IndexError from None
743                in_indices[i] = indx
744            elif indx.dtype.kind != "b" and indx.dtype.kind != "i":
745                raise IndexError(
746                    "arrays used as indices must be of integer (or boolean) type"
747                )
748            if indx.ndim != 0:
749                no_copy = False
750            ndim += 1
751            fancy_dim += 1
752
753        if arr.ndim - ndim < 0:
754            # we can't take more dimensions then we have, not even for 0-d
755            # arrays.  since a[()] makes sense, but not a[(),]. We will
756            # raise an error later on, unless a broadcasting error occurs
757            # first.
758            raise IndexError
759
760        if ndim == 0 and None not in in_indices:
761            # Well we have no indexes or one Ellipsis. This is legal.
762            return arr.copy(), no_copy
763
764        if ellipsis_pos is not None:
765            in_indices[ellipsis_pos : ellipsis_pos + 1] = [slice(None, None)] * (
766                arr.ndim - ndim
767            )
768
769        for ax, indx in enumerate(in_indices):
770            if isinstance(indx, slice):
771                # convert to an index array
772                indx = np.arange(*indx.indices(arr.shape[ax]))
773                indices.append(["s", indx])
774                continue
775            elif indx is None:
776                # this is like taking a slice with one element from a new axis:
777                indices.append(["n", np.array([0], dtype=np.intp)])
778                arr = arr.reshape(arr.shape[:ax] + (1,) + arr.shape[ax:])
779                continue
780            if isinstance(indx, np.ndarray) and indx.dtype == bool:
781                if indx.shape != arr.shape[ax : ax + indx.ndim]:
782                    raise IndexError
783
784                try:
785                    flat_indx = np.ravel_multi_index(
786                        np.nonzero(indx), arr.shape[ax : ax + indx.ndim], mode="raise"
787                    )
788                except Exception:
789                    error_unless_broadcast_to_empty = True
790                    # fill with 0s instead, and raise error later
791                    flat_indx = np.array([0] * indx.sum(), dtype=np.intp)
792                # concatenate axis into a single one:
793                if indx.ndim != 0:
794                    arr = arr.reshape(
795                        arr.shape[:ax]
796                        + (np.prod(arr.shape[ax : ax + indx.ndim]),)
797                        + arr.shape[ax + indx.ndim :]
798                    )
799                    indx = flat_indx
800                else:
801                    # This could be changed, a 0-d boolean index can
802                    # make sense (even outside the 0-d indexed array case)
803                    # Note that originally this is could be interpreted as
804                    # integer in the full integer special case.
805                    raise IndexError
806            else:
807                # If the index is a singleton, the bounds check is done
808                # before the broadcasting. This used to be different in <1.9
809                if indx.ndim == 0:
810                    if indx >= arr.shape[ax] or indx < -arr.shape[ax]:
811                        raise IndexError
812            if indx.ndim == 0:
813                # The index is a scalar. This used to be two fold, but if
814                # fancy indexing was active, the check was done later,
815                # possibly after broadcasting it away (1.7. or earlier).
816                # Now it is always done.
817                if indx >= arr.shape[ax] or indx < -arr.shape[ax]:
818                    raise IndexError
819            if len(indices) > 0 and indices[-1][0] == "f" and ax != ellipsis_pos:
820                # NOTE: There could still have been a 0-sized Ellipsis
821                # between them. Checked that with ellipsis_pos.
822                indices[-1].append(indx)
823            else:
824                # We have a fancy index that is not after an existing one.
825                # NOTE: A 0-d array triggers this as well, while one may
826                # expect it to not trigger it, since a scalar would not be
827                # considered fancy indexing.
828                num_fancy += 1
829                indices.append(["f", indx])
830
831        if num_fancy > 1 and not no_copy:
832            # We have to flush the fancy indexes left
833            new_indices = indices[:]
834            axes = list(range(arr.ndim))
835            fancy_axes = []
836            new_indices.insert(0, ["f"])
837            ni = 0
838            ai = 0
839            for indx in indices:
840                ni += 1
841                if indx[0] == "f":
842                    new_indices[0].extend(indx[1:])
843                    del new_indices[ni]
844                    ni -= 1
845                    for ax in range(ai, ai + len(indx[1:])):
846                        fancy_axes.append(ax)
847                        axes.remove(ax)
848                ai += len(indx) - 1  # axis we are at
849            indices = new_indices
850            # and now we need to transpose arr:
851            arr = arr.transpose(*(fancy_axes + axes))
852
853        # We only have one 'f' index now and arr is transposed accordingly.
854        # Now handle newaxis by reshaping...
855        ax = 0
856        for indx in indices:
857            if indx[0] == "f":
858                if len(indx) == 1:
859                    continue
860                # First of all, reshape arr to combine fancy axes into one:
861                orig_shape = arr.shape
862                orig_slice = orig_shape[ax : ax + len(indx[1:])]
863                arr = arr.reshape(
864                    arr.shape[:ax]
865                    + (np.prod(orig_slice).astype(int),)
866                    + arr.shape[ax + len(indx[1:]) :]
867                )
868
869                # Check if broadcasting works
870                res = np.broadcast(*indx[1:])
871                # unfortunately the indices might be out of bounds. So check
872                # that first, and use mode='wrap' then. However only if
873                # there are any indices...
874                if res.size != 0:
875                    if error_unless_broadcast_to_empty:
876                        raise IndexError
877                    for _indx, _size in zip(indx[1:], orig_slice):
878                        if _indx.size == 0:
879                            continue
880                        if np.any(_indx >= _size) or np.any(_indx < -_size):
881                            raise IndexError
882                if len(indx[1:]) == len(orig_slice):
883                    if np.prod(orig_slice) == 0:
884                        # Work around for a crash or IndexError with 'wrap'
885                        # in some 0-sized cases.
886                        try:
887                            mi = np.ravel_multi_index(
888                                indx[1:], orig_slice, mode="raise"
889                            )
890                        except Exception as exc:
891                            # This happens with 0-sized orig_slice (sometimes?)
892                            # here it is a ValueError, but indexing gives a:
893                            raise IndexError("invalid index into 0-sized") from exc
894                    else:
895                        mi = np.ravel_multi_index(indx[1:], orig_slice, mode="wrap")
896                else:
897                    # Maybe never happens...
898                    raise ValueError
899                arr = arr.take(mi.ravel(), axis=ax)
900                try:
901                    arr = arr.reshape(arr.shape[:ax] + mi.shape + arr.shape[ax + 1 :])
902                except ValueError:
903                    # too many dimensions, probably
904                    raise IndexError from None
905                ax += mi.ndim
906                continue
907
908            # If we are here, we have a 1D array for take:
909            arr = arr.take(indx[1], axis=ax)
910            ax += 1
911
912        return arr, no_copy
913
914    def _check_multi_index(self, arr, index):
915        """Check a multi index item getting and simple setting.
916
917        Parameters
918        ----------
919        arr : ndarray
920            Array to be indexed, must be a reshaped arange.
921        index : tuple of indexing objects
922            Index being tested.
923        """
924        # Test item getting
925        try:
926            mimic_get, no_copy = self._get_multi_index(arr, index)
927        except Exception as e:
928            if HAS_REFCOUNT:
929                prev_refcount = sys.getrefcount(arr)
930            assert_raises(type(e), arr.__getitem__, index)
931            assert_raises(type(e), arr.__setitem__, index, 0)
932            if HAS_REFCOUNT:
933                assert_equal(prev_refcount, sys.getrefcount(arr))
934            return
935
936        self._compare_index_result(arr, index, mimic_get, no_copy)
937
938    def _check_single_index(self, arr, index):
939        """Check a single index item getting and simple setting.
940
941        Parameters
942        ----------
943        arr : ndarray
944            Array to be indexed, must be an arange.
945        index : indexing object
946            Index being tested. Must be a single index and not a tuple
947            of indexing objects (see also `_check_multi_index`).
948        """
949        try:
950            mimic_get, no_copy = self._get_multi_index(arr, (index,))
951        except Exception as e:
952            if HAS_REFCOUNT:
953                prev_refcount = sys.getrefcount(arr)
954            assert_raises(type(e), arr.__getitem__, index)
955            assert_raises(type(e), arr.__setitem__, index, 0)
956            if HAS_REFCOUNT:
957                assert_equal(prev_refcount, sys.getrefcount(arr))
958            return
959
960        self._compare_index_result(arr, index, mimic_get, no_copy)
961
962    def _compare_index_result(self, arr, index, mimic_get, no_copy):
963        """Compare mimicked result to indexing result."""
964        raise SkipTest("torch does not support subclassing")
965        arr = arr.copy()
966        indexed_arr = arr[index]
967        assert_array_equal(indexed_arr, mimic_get)
968        # Check if we got a view, unless its a 0-sized or 0-d array.
969        # (then its not a view, and that does not matter)
970        if indexed_arr.size != 0 and indexed_arr.ndim != 0:
971            assert_(np.may_share_memory(indexed_arr, arr) == no_copy)
972            # Check reference count of the original array
973            if HAS_REFCOUNT:
974                if no_copy:
975                    # refcount increases by one:
976                    assert_equal(sys.getrefcount(arr), 3)
977                else:
978                    assert_equal(sys.getrefcount(arr), 2)
979
980        # Test non-broadcast setitem:
981        b = arr.copy()
982        b[index] = mimic_get + 1000
983        if b.size == 0:
984            return  # nothing to compare here...
985        if no_copy and indexed_arr.ndim != 0:
986            # change indexed_arr in-place to manipulate original:
987            indexed_arr += 1000
988            assert_array_equal(arr, b)
989            return
990        # Use the fact that the array is originally an arange:
991        arr.flat[indexed_arr.ravel()] += 1000
992        assert_array_equal(arr, b)
993
994    def test_boolean(self):
995        a = np.array(5)
996        assert_equal(a[np.array(True)], 5)
997        a[np.array(True)] = 1
998        assert_equal(a, 1)
999        # NOTE: This is different from normal broadcasting, as
1000        # arr[boolean_array] works like in a multi index. Which means
1001        # it is aligned to the left. This is probably correct for
1002        # consistency with arr[boolean_array,] also no broadcasting
1003        # is done at all
1004        self._check_multi_index(self.a, (np.zeros_like(self.a, dtype=bool),))
1005        self._check_multi_index(self.a, (np.zeros_like(self.a, dtype=bool)[..., 0],))
1006        self._check_multi_index(self.a, (np.zeros_like(self.a, dtype=bool)[None, ...],))
1007
1008    def test_multidim(self):
1009        # Automatically test combinations with complex indexes on 2nd (or 1st)
1010        # spot and the simple ones in one other spot.
1011        with warnings.catch_warnings():
1012            # This is so that np.array(True) is not accepted in a full integer
1013            # index, when running the file separately.
1014            warnings.filterwarnings("error", "", DeprecationWarning)
1015            warnings.filterwarnings("error", "", np.VisibleDeprecationWarning)
1016
1017            def isskip(idx):
1018                return isinstance(idx, str) and idx == "skip"
1019
1020            for simple_pos in [0, 2, 3]:
1021                tocheck = [
1022                    self.fill_indices,
1023                    self.complex_indices,
1024                    self.fill_indices,
1025                    self.fill_indices,
1026                ]
1027                tocheck[simple_pos] = self.simple_indices
1028                for index in product(*tocheck):
1029                    index = tuple(i for i in index if not isskip(i))
1030                    self._check_multi_index(self.a, index)
1031                    self._check_multi_index(self.b, index)
1032
1033        # Check very simple item getting:
1034        self._check_multi_index(self.a, (0, 0, 0, 0))
1035        self._check_multi_index(self.b, (0, 0, 0, 0))
1036        # Also check (simple cases of) too many indices:
1037        assert_raises(IndexError, self.a.__getitem__, (0, 0, 0, 0, 0))
1038        assert_raises(IndexError, self.a.__setitem__, (0, 0, 0, 0, 0), 0)
1039        assert_raises(IndexError, self.a.__getitem__, (0, 0, [1], 0, 0))
1040        assert_raises(IndexError, self.a.__setitem__, (0, 0, [1], 0, 0), 0)
1041
1042    def test_1d(self):
1043        a = np.arange(10)
1044        for index in self.complex_indices:
1045            self._check_single_index(a, index)
1046
1047
1048class TestFloatNonIntegerArgument(TestCase):
1049    """
1050    These test that ``TypeError`` is raised when you try to use
1051    non-integers as arguments to for indexing and slicing e.g. ``a[0.0:5]``
1052    and ``a[0.5]``, or other functions like ``array.reshape(1., -1)``.
1053
1054    """
1055
1056    def test_valid_indexing(self):
1057        # These should raise no errors.
1058        a = np.array([[[5]]])
1059
1060        a[np.array([0])]
1061        a[[0, 0]]
1062        a[:, [0, 0]]
1063        a[:, 0, :]
1064        a[:, :, :]
1065
1066    def test_valid_slicing(self):
1067        # These should raise no errors.
1068        a = np.array([[[5]]])
1069
1070        a[::]
1071        a[0:]
1072        a[:2]
1073        a[0:2]
1074        a[::2]
1075        a[1::2]
1076        a[:2:2]
1077        a[1:2:2]
1078
1079    def test_non_integer_argument_errors(self):
1080        a = np.array([[5]])
1081
1082        assert_raises(TypeError, np.reshape, a, (1.0, 1.0, -1))
1083        assert_raises(TypeError, np.reshape, a, (np.array(1.0), -1))
1084        assert_raises(TypeError, np.take, a, [0], 1.0)
1085        assert_raises((TypeError, RuntimeError), np.take, a, [0], np.float64(1.0))
1086
1087    @skip(
1088        reason=("torch doesn't have scalar types with distinct element-wise behaviours")
1089    )
1090    def test_non_integer_sequence_multiplication(self):
1091        # NumPy scalar sequence multiply should not work with non-integers
1092        def mult(a, b):
1093            return a * b
1094
1095        assert_raises(TypeError, mult, [1], np.float64(3))
1096        # following should be OK
1097        mult([1], np.int_(3))
1098
1099    def test_reduce_axis_float_index(self):
1100        d = np.zeros((3, 3, 3))
1101        assert_raises(TypeError, np.min, d, 0.5)
1102        assert_raises(TypeError, np.min, d, (0.5, 1))
1103        assert_raises(TypeError, np.min, d, (1, 2.2))
1104        assert_raises(TypeError, np.min, d, (0.2, 1.2))
1105
1106
1107class TestBooleanIndexing(TestCase):
1108    # Using a boolean as integer argument/indexing is an error.
1109    def test_bool_as_int_argument_errors(self):
1110        a = np.array([[[1]]])
1111
1112        assert_raises(TypeError, np.reshape, a, (True, -1))
1113        # Note that operator.index(np.array(True)) does not work, a boolean
1114        # array is thus also deprecated, but not with the same message:
1115        # assert_warns(DeprecationWarning, operator.index, np.True_)
1116
1117        assert_raises(TypeError, np.take, args=(a, [0], False))
1118
1119        raise SkipTest("torch consumes boolean tensors as ints, no bother raising here")
1120        assert_raises(TypeError, np.reshape, a, (np.bool_(True), -1))
1121        assert_raises(TypeError, operator.index, np.array(True))
1122
1123    def test_boolean_indexing_weirdness(self):
1124        # Weird boolean indexing things
1125        a = np.ones((2, 3, 4))
1126        assert a[False, True, ...].shape == (0, 2, 3, 4)
1127        assert a[True, [0, 1], True, True, [1], [[2]]].shape == (1, 2)
1128        assert_raises(IndexError, lambda: a[False, [0, 1], ...])
1129
1130    def test_boolean_indexing_fast_path(self):
1131        # These used to either give the wrong error, or incorrectly give no
1132        # error.
1133        a = np.ones((3, 3))
1134
1135        # This used to incorrectly work (and give an array of shape (0,))
1136        idx1 = np.array([[False] * 9])
1137        with pytest.raises(IndexError):
1138            a[idx1]
1139
1140        # This used to incorrectly give a ValueError: operands could not be broadcast together
1141        idx2 = np.array([[False] * 8 + [True]])
1142        with pytest.raises(IndexError):
1143            a[idx2]
1144
1145        # This is the same as it used to be. The above two should work like this.
1146        idx3 = np.array([[False] * 10])
1147        with pytest.raises(IndexError):
1148            a[idx3]
1149
1150        # This used to give ValueError: non-broadcastable operand
1151        a = np.ones((1, 1, 2))
1152        idx = np.array([[[True], [False]]])
1153        with pytest.raises(IndexError):
1154            a[idx]
1155
1156
1157class TestArrayToIndexDeprecation(TestCase):
1158    """Creating an index from array not 0-D is an error."""
1159
1160    def test_array_to_index_error(self):
1161        # so no exception is expected. The raising is effectively tested above.
1162        a = np.array([[[1]]])
1163
1164        assert_raises((TypeError, RuntimeError), np.take, a, [0], a)
1165
1166        raise SkipTest(
1167            "Multi-dimensional tensors are indexable just as long as they only "
1168            "contain a single element, no bother raising here"
1169        )
1170        assert_raises(TypeError, operator.index, np.array([1]))
1171
1172        raise SkipTest("torch consumes tensors as ints, no bother raising here")
1173        assert_raises(TypeError, np.reshape, a, (a, -1))
1174
1175
1176class TestNonIntegerArrayLike(TestCase):
1177    """Tests that array_likes only valid if can safely cast to integer.
1178
1179    For instance, lists give IndexError when they cannot be safely cast to
1180    an integer.
1181
1182    """
1183
1184    @skip(
1185        reason=(
1186            "torch consumes floats by way of falling back on its deprecated "
1187            "__index__ behaviour, no bother raising here"
1188        )
1189    )
1190    def test_basic(self):
1191        a = np.arange(10)
1192
1193        assert_raises(IndexError, a.__getitem__, [0.5, 1.5])
1194        assert_raises(IndexError, a.__getitem__, (["1", "2"],))
1195
1196        # The following is valid
1197        a.__getitem__([])
1198
1199
1200class TestMultipleEllipsisError(TestCase):
1201    """An index can only have a single ellipsis."""
1202
1203    @xfail  # (
1204    #    reason=(
1205    #        "torch currently consumes multiple ellipsis, no bother raising "
1206    #        "here. See https://github.com/pytorch/pytorch/issues/59787#issue-917252204"
1207    #    )
1208    # )
1209    def test_basic(self):
1210        a = np.arange(10)
1211        assert_raises(IndexError, lambda: a[..., ...])
1212        assert_raises(IndexError, a.__getitem__, ((Ellipsis,) * 2,))
1213        assert_raises(IndexError, a.__getitem__, ((Ellipsis,) * 3,))
1214
1215
1216if __name__ == "__main__":
1217    run_tests()
1218