From e1dc133ac0b5f6eb025640d64b2f7174425fe713 Mon Sep 17 00:00:00 2001 From: Rishi Kulkarni Date: Tue, 21 Dec 2021 17:22:13 -0500 Subject: [PATCH 1/2] performance improvements to np.full and np.ones --- numba/np/arrayobj.py | 40 ++++++++++++++++++++++++---------------- 1 file changed, 24 insertions(+), 16 deletions(-) diff --git a/numba/np/arrayobj.py b/numba/np/arrayobj.py index a7046a404f4..6722193bb83 100644 --- a/numba/np/arrayobj.py +++ b/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 enumerate(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 enumerate(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 enumerate(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 enumerate(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 enumerate(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 enumerate(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 enumerate(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 enumerate(arr_flat): + arr_flat[idx] = 1 return arr res = context.compile_internal(builder, ones_like, sig, args) From 540c9100eeea2c0422e8ccf7493fac3fae89bed4 Mon Sep 17 00:00:00 2001 From: Rishi Kulkarni Date: Thu, 20 Jan 2022 09:38:46 -0500 Subject: [PATCH 2/2] use range(len()) instead of enumerate --- numba/np/arrayobj.py | 16 ++++++++-------- 1 file changed, 8 insertions(+), 8 deletions(-) diff --git a/numba/np/arrayobj.py b/numba/np/arrayobj.py index 6722193bb83..b655753e61e 100644 --- a/numba/np/arrayobj.py +++ b/numba/np/arrayobj.py @@ -3790,7 +3790,7 @@ def numpy_full_nd(context, builder, sig, args): def full(shape, value): arr = np.empty(shape, type(value)) arr_flat = arr.flat - for idx, _ in enumerate(arr_flat): + for idx in range(len(arr_flat)): arr_flat[idx] = value return arr @@ -3805,7 +3805,7 @@ def numpy_full_dtype_nd(context, builder, sig, args): def full(shape, value, dtype): arr = np.empty(shape, dtype) arr_flat = arr.flat - for idx, _ in enumerate(arr_flat): + for idx in range(len(arr_flat)): arr_flat[idx] = value return arr @@ -3819,7 +3819,7 @@ def numpy_full_like_nd(context, builder, sig, args): def full_like(arr, value): arr = np.empty_like(arr) arr_flat = arr.flat - for idx, _ in enumerate(arr_flat): + for idx in range(len(arr_flat)): arr_flat[idx] = value return arr @@ -3834,7 +3834,7 @@ def numpy_full_like_nd_type_spec(context, builder, sig, args): def full_like(arr, value, dtype): arr = np.empty_like(arr, dtype) arr_flat = arr.flat - for idx, _ in enumerate(arr_flat): + for idx in range(len(arr_flat)): arr_flat[idx] = value return arr @@ -3848,7 +3848,7 @@ def numpy_ones_nd(context, builder, sig, args): def ones(shape): arr = np.empty(shape) arr_flat = arr.flat - for idx, _ in enumerate(arr_flat): + for idx in range(len(arr_flat)): arr_flat[idx] = 1 return arr @@ -3865,7 +3865,7 @@ def numpy_ones_dtype_nd(context, builder, sig, args): def ones(shape, dtype): arr = np.empty(shape, dtype) arr_flat = arr.flat - for idx, _ in enumerate(arr_flat): + for idx in range(len(arr_flat)): arr_flat[idx] = 1 return arr @@ -3879,7 +3879,7 @@ def numpy_ones_like_nd(context, builder, sig, args): def ones_like(arr): arr = np.empty_like(arr) arr_flat = arr.flat - for idx, _ in enumerate(arr_flat): + for idx in range(len(arr_flat)): arr_flat[idx] = 1 return arr @@ -3894,7 +3894,7 @@ def numpy_ones_like_dtype_nd(context, builder, sig, args): def ones_like(arr, dtype): arr = np.empty_like(arr, dtype) arr_flat = arr.flat - for idx, _ in enumerate(arr_flat): + for idx in range(len(arr_flat)): arr_flat[idx] = 1 return arr