Repository: numba/numba

Issue:
An error with ndarray fancy indexing with ellipsis
The following code demonstrates the problem

``` python
from numba import njit
import numpy as np

@njit
def ellipsis_error(arr):
    out = np.zeros_like(arr)
    i = 0
    for index, val in np.ndenumerate(arr[..., i]):  # can be "fixed" by using colon `:` instead of `...`
        out[index][i] = val

    return out

py_func = ellipsis_error.py_func

outersize = 4
innersize = 4
arr = np.arange(outersize * innersize).reshape(outersize, innersize)
print('---arr---')
print(arr)
got = ellipsis_error(arr)
expected = py_func(arr)

print('---got---')
print(got)

print('---expected---')
print(expected)

np.testing.assert_equal(got, expected)
```

Prints

``` python
---arr---
[[ 0  1  2  3]
 [ 4  5  6  7]
 [ 8  9 10 11]
 [12 13 14 15]]
---got---
[[0 0 0 0]
 [1 0 0 0]
 [2 0 0 0]
 [3 0 0 0]]
---expected---
[[ 0  0  0  0]
 [ 4  0  0  0]
 [ 8  0  0  0]
 [12  0  0  0]]
```

Note the error is likely to be error in strides computation.

An error with ndarray fancy indexing with ellipsis
The following code demonstrates the problem

``` python
from numba import njit
import numpy as np

@njit
def ellipsis_error(arr):
    out = np.zeros_like(arr)
    i = 0
    for index, val in np.ndenumerate(arr[..., i]):  # can be "fixed" by using colon `:` instead of `...`
        out[index][i] = val

    return out

py_func = ellipsis_error.py_func

outersize = 4
innersize = 4
arr = np.arange(outersize * innersize).reshape(outersize, innersize)
print('---arr---')
print(arr)
got = ellipsis_error(arr)
expected = py_func(arr)

print('---got---')
print(got)

print('---expected---')
print(expected)

np.testing.assert_equal(got, expected)
```

Prints

``` python
---arr---
[[ 0  1  2  3]
 [ 4  5  6  7]
 [ 8  9 10 11]
 [12 13 14 15]]
---got---
[[0 0 0 0]
 [1 0 0 0]
 [2 0 0 0]
 [3 0 0 0]]
---expected---
[[ 0  0  0  0]
 [ 4  0  0  0]
 [ 8  0  0  0]
 [12  0  0  0]]
```

Note the error is likely to be error in strides computation.



Discussion:

Resolving patch:
diff --git a/numba/typing/arraydecl.py b/numba/typing/arraydecl.py
--- a/numba/typing/arraydecl.py
+++ b/numba/typing/arraydecl.py
@@ -96,7 +96,7 @@ def get_array_index_type(ary, idx):
         layout = ary.layout
         if layout == 'C':
             # Integer indexing on the left keeps the array C-contiguous
-            if len(left_indices) + len(right_indices) == ary.ndim:
+            if n_indices == ary.ndim:
                 # If all indices are there, ellipsis's place is indifferent
                 left_indices = left_indices + right_indices
                 right_indices = []
@@ -110,7 +110,7 @@ def get_array_index_type(ary, idx):
                         break
         elif layout == 'F':
             # Integer indexing on the right keeps the array F-contiguous
-            if len(left_indices) + len(right_indices) == ary.ndim:
+            if n_indices == ary.ndim:
                 # If all indices are there, ellipsis's place is indifferent
                 right_indices = left_indices + right_indices
                 left_indices = []


Test patch:
diff --git a/numba/tests/test_indexing.py b/numba/tests/test_indexing.py
--- a/numba/tests/test_indexing.py
+++ b/numba/tests/test_indexing.py
@@ -630,6 +630,27 @@ def test_ellipsis3(self, flags=enable_pyobj_flags):
     def test_ellipsis3_npm(self):
         self.test_ellipsis3(flags=Noflags)
 
+    def test_ellipsis_issue1498(self):
+        # This is an issue due to incorrect layout inferred for when
+        # ellpsis is used and ndenumerate is specializing on the layout.
+        @njit
+        def udt(arr):
+            out = np.zeros_like(arr)
+            i = 0
+            for index, val in np.ndenumerate(arr[..., i]):
+                out[index][i] = val
+
+            return out
+
+        py_func = udt.py_func
+
+        outersize = 4
+        innersize = 4
+        arr = np.arange(outersize * innersize).reshape(outersize, innersize)
+        got = udt(arr)
+        expected = py_func(arr)
+        np.testing.assert_equal(got, expected)
+
     def test_none_index(self, flags=enable_pyobj_flags):
         pyfunc = none_index_usecase
         arraytype = types.Array(types.int32, 2, 'C')