Repository: numba/numba

Issue:
Yielding a ndindex raises an error in nopython mode
``` python
import numba
import numpy as np

@numba.jit(nopython=True)
def f():
    for ind in np.ndindex(5):
        yield ind[0]

for i in f():
    print(i)
```

This code raises the following error: `KeyError: "Failed at nopython (nopython mode backend)\n<class 'numba.types.NumpyNdIndexType'>"`.

It works in nopython=False mode.

Numba 0.18.2
Python 3.4.3, Windows 7 64 bits



Discussion:
It compiles here but the generator never finishes.


Resolving patch:
diff --git a/numba/datamodel/models.py b/numba/datamodel/models.py
--- a/numba/datamodel/models.py
+++ b/numba/datamodel/models.py
@@ -254,6 +254,24 @@ def load_from_data_pointer(self, builder, ptr):
         return builder.bitcast(ptr, self.get_value_type())
 
 
+@register_default(types.EphemeralArray)
+class EphemeralArrayModel(PointerModel):
+
+    def get_data_type(self):
+        return ir.ArrayType(self._pointee_be_type, self._fe_type.count)
+
+    def as_data(self, builder, value):
+        values = [builder.load(cgutils.gep(builder, value, i))
+                  for i in range(self._fe_type.count)]
+        return cgutils.pack_array(builder, values)
+
+    def from_data(self, builder, value):
+        raise NotImplementedError("use load_from_data_pointer() instead")
+
+    def load_from_data_pointer(self, builder, ptr):
+        return builder.bitcast(ptr, self.get_value_type())
+
+
 @register_default(types.ExternalFunctionPointer)
 class ExternalFuncPointerModel(PrimitiveModel):
     def __init__(self, dmm, fe_type):
@@ -686,40 +704,34 @@ def from_argument(self, builder, value):
 
 
 class CContiguousFlatIter(StructModel):
-    def __init__(self, dmm, fe_type):
+    def __init__(self, dmm, fe_type, need_indices):
         assert fe_type.array_type.layout == 'C'
         array_type = fe_type.array_type
         dtype = array_type.dtype
-        members = [('array', types.CPointer(array_type)),
+        ndim = array_type.ndim
+        members = [('array', types.EphemeralPointer(array_type)),
                    ('stride', types.intp),
                    ('pointer', types.EphemeralPointer(types.CPointer(dtype))),
                    ('index', types.EphemeralPointer(types.intp)),
-                   # NOTE: indices is an array
-                   ('indices', types.EphemeralPointer(types.intp)),
-        ]
+                   ]
+        if need_indices:
+            # For ndenumerate()
+            members.append(('indices', types.EphemeralArray(types.intp, ndim)))
         super(CContiguousFlatIter, self).__init__(dmm, fe_type, members)
 
-    def get_data_type(self):
-        # FIXME: must flatten indices array
-        raise NotImplementedError("cannot serialize flat() iterator")
-
 
 class FlatIter(StructModel):
     def __init__(self, dmm, fe_type):
         array_type = fe_type.array_type
         dtype = array_type.dtype
-        members = [('array', types.CPointer(array_type)),
-                   # NOTE: pointers and indices are arrays
-                   ('pointers', types.EphemeralPointer(types.CPointer(dtype))),
-                   ('indices', types.EphemeralPointer(types.intp)),
+        ndim = array_type.ndim
+        members = [('array', types.EphemeralPointer(array_type)),
+                   ('pointers', types.EphemeralArray(types.CPointer(dtype), ndim)),
+                   ('indices', types.EphemeralArray(types.intp, ndim)),
                    ('exhausted', types.EphemeralPointer(types.boolean)),
         ]
         super(FlatIter, self).__init__(dmm, fe_type, members)
 
-    def get_data_type(self):
-        # FIXME: must flatten indices and pointers
-        raise NotImplementedError("cannot serialize flat() iterator")
-
 
 @register_default(types.UniTupleIter)
 class UniTupleIter(StructModel):
@@ -852,17 +864,23 @@ class NdIndexType(StructModel):
     def __init__(self, dmm, fe_type):
         ndim = fe_type.ndim
         members = [('shape', types.UniTuple(types.intp, ndim)),
-                   ('indices', types.CPointer(types.intp)),
-                   ('exhausted', types.CPointer(types.boolean)),
+                   ('indices', types.EphemeralArray(types.intp, ndim)),
+                   ('exhausted', types.EphemeralPointer(types.boolean)),
                    ]
         super(NdIndexType, self).__init__(dmm, fe_type, members)
 
 
 @register_default(types.NumpyFlatType)
-@register_default(types.NumpyNdEnumerateType)
 def handle_numpy_flat_type(dmm, ty):
     if ty.array_type.layout == 'C':
-        return CContiguousFlatIter(dmm, ty)
+        return CContiguousFlatIter(dmm, ty, need_indices=False)
+    else:
+        return FlatIter(dmm, ty)
+
+@register_default(types.NumpyNdEnumerateType)
+def handle_numpy_ndenumerate_type(dmm, ty):
+    if ty.array_type.layout == 'C':
+        return CContiguousFlatIter(dmm, ty, need_indices=True)
     else:
         return FlatIter(dmm, ty)
 
diff --git a/numba/types.py b/numba/types.py
--- a/numba/types.py
+++ b/numba/types.py
@@ -1082,6 +1082,23 @@ class EphemeralPointer(CPointer):
     """
 
 
+class EphemeralArray(Type):
+    """
+    Similar to EphemeralPointer, but pointing to an array of elements,
+    rather than a single one.  The array size must be known at compile-time.
+    """
+
+    def __init__(self, dtype, count):
+        self.dtype = dtype
+        self.count = count
+        name = "*%s[%d]" % (dtype, count)
+        super(EphemeralArray, self).__init__(name, param=True)
+
+    @property
+    def key(self):
+        return self.dtype, self.count
+
+
 class Object(Type):
     mutable = True
 


Test patch:
diff --git a/numba/tests/test_generators.py b/numba/tests/test_generators.py
--- a/numba/tests/test_generators.py
+++ b/numba/tests/test_generators.py
@@ -80,6 +80,19 @@ def return_generator_expr(x):
     return (i*2 for i in x)
 
 
+def gen_ndindex(shape):
+    for ind in np.ndindex(shape):
+        yield ind
+
+def gen_flat(arr):
+    for val in arr.flat:
+        yield val
+
+def gen_ndenumerate(arr):
+    for tup in np.ndenumerate(arr):
+        yield tup
+
+
 class TestGenerators(TestCase):
 
     def check_generator(self, pygen, cgen):
@@ -204,6 +217,45 @@ def test_consume_gen2(self):
     def test_consume_gen3(self):
         self.check_consume_generator(gen3)
 
+    # Check generator storage of some types
+
+    def check_ndindex(self, flags=no_pyobj_flags):
+        pyfunc = gen_ndindex
+        cr = compile_isolated(pyfunc, (types.UniTuple(types.intp, 2),),
+                              flags=flags)
+        shape = (2, 3)
+        pygen = pyfunc(shape)
+        cgen = cr.entry_point(shape)
+        self.check_generator(pygen, cgen)
+
+    def test_ndindex(self):
+        self.check_ndindex()
+
+    def test_ndindex_objmode(self):
+        self.check_ndindex(flags=forceobj_flags)
+
+    def check_np_flat(self, pyfunc, flags=no_pyobj_flags):
+        cr = compile_isolated(pyfunc, (types.Array(types.intp, 2, "C"),),
+                              flags=flags)
+        arr = np.arange(6).reshape((2, 3))
+        self.check_generator(pyfunc(arr), cr.entry_point(arr))
+        cr = compile_isolated(pyfunc, (types.Array(types.intp, 2, "A"),),
+                              flags=flags)
+        arr = arr.T
+        self.check_generator(pyfunc(arr), cr.entry_point(arr))
+
+    def test_np_flat(self):
+        self.check_np_flat(gen_flat)
+
+    def test_np_flat_objmode(self):
+        self.check_np_flat(gen_flat, flags=forceobj_flags)
+
+    def test_ndenumerate(self):
+        self.check_np_flat(gen_ndenumerate)
+
+    def test_ndenumerate_objmode(self):
+        self.check_np_flat(gen_ndenumerate, flags=forceobj_flags)
+
 
 class TestGenExprs(TestCase):