Merge pull request numba#7682 from rishi-kulkarni/master

performance improvements to np.full and np.ones

numba/np/arrayobj.py

@@ -3789,8 +3789,9 @@ def numpy_full_nd(context, builder, sig, args):
 
     def full(shape, value):
         arr = np.empty(shape, type(value))
-        for idx in np.ndindex(arr.shape):
-            arr[idx] = value
+        arr_flat = arr.flat
+        for idx in range(len(arr_flat)):
+            arr_flat[idx] = value
         return arr
 
     res = context.compile_internal(builder, full, sig, args)
@@ -3803,8 +3804,9 @@ def numpy_full_dtype_nd(context, builder, sig, args):
 
     def full(shape, value, dtype):
         arr = np.empty(shape, dtype)
-        for idx in np.ndindex(arr.shape):
-            arr[idx] = value
+        arr_flat = arr.flat
+        for idx in range(len(arr_flat)):
+            arr_flat[idx] = value
         return arr
 
     res = context.compile_internal(builder, full, sig, args)
@@ -3816,8 +3818,9 @@ def numpy_full_like_nd(context, builder, sig, args):
 
     def full_like(arr, value):
         arr = np.empty_like(arr)
-        for idx in np.ndindex(arr.shape):
-            arr[idx] = value
+        arr_flat = arr.flat
+        for idx in range(len(arr_flat)):
+            arr_flat[idx] = value
         return arr
 
     res = context.compile_internal(builder, full_like, sig, args)
@@ -3830,8 +3833,9 @@ def numpy_full_like_nd_type_spec(context, builder, sig, args):
 
     def full_like(arr, value, dtype):
         arr = np.empty_like(arr, dtype)
-        for idx in np.ndindex(arr.shape):
-            arr[idx] = value
+        arr_flat = arr.flat
+        for idx in range(len(arr_flat)):
+            arr_flat[idx] = value
         return arr
 
     res = context.compile_internal(builder, full_like, sig, args)
@@ -3843,8 +3847,9 @@ def numpy_ones_nd(context, builder, sig, args):
 
     def ones(shape):
         arr = np.empty(shape)
-        for idx in np.ndindex(arr.shape):
-            arr[idx] = 1
+        arr_flat = arr.flat
+        for idx in range(len(arr_flat)):
+            arr_flat[idx] = 1
         return arr
 
     valty = sig.return_type.dtype
@@ -3859,8 +3864,9 @@ def numpy_ones_dtype_nd(context, builder, sig, args):
 
     def ones(shape, dtype):
         arr = np.empty(shape, dtype)
-        for idx in np.ndindex(arr.shape):
-            arr[idx] = 1
+        arr_flat = arr.flat
+        for idx in range(len(arr_flat)):
+            arr_flat[idx] = 1
         return arr
 
     res = context.compile_internal(builder, ones, sig, args)
@@ -3872,8 +3878,9 @@ def numpy_ones_like_nd(context, builder, sig, args):
 
     def ones_like(arr):
         arr = np.empty_like(arr)
-        for idx in np.ndindex(arr.shape):
-            arr[idx] = 1
+        arr_flat = arr.flat
+        for idx in range(len(arr_flat)):
+            arr_flat[idx] = 1
         return arr
 
     res = context.compile_internal(builder, ones_like, sig, args)
@@ -3886,8 +3893,9 @@ def numpy_ones_like_dtype_nd(context, builder, sig, args):
 
     def ones_like(arr, dtype):
         arr = np.empty_like(arr, dtype)
-        for idx in np.ndindex(arr.shape):
-            arr[idx] = 1
+        arr_flat = arr.flat
+        for idx in range(len(arr_flat)):
+            arr_flat[idx] = 1
         return arr
 
     res = context.compile_internal(builder, ones_like, sig, args)