Optimize concatenate when number of array is small
@@ -2094,13 +2094,15 @@ cpdef ndarray concatenate(tup, axis, shape, dtype):
ret = ndarray(shape, dtype=dtype)
+ if len(tup) > 3:
all_same_type = True
all_one_and_contiguous = True
dtype = tup[0].dtype
for a in tup:
all_same_type = all_same_type and (a.dtype == dtype)
all_one_and_contiguous = (
- all_one_and_contiguous and a._c_contiguous and a._shape[axis] == 1)
+ all_one_and_contiguous and a._c_contiguous and
+ a._shape[axis] == 1)
if all_same_type:
x = array([a.data.ptr for a in tup])
@@ -2118,7 +2120,8 @@ cpdef ndarray concatenate(tup, axis, shape, dtype):
_concatenate_kernel(
x, axis, len(shape), cum_sizes, x_strides, ret)
- else:
+ return ret
+
skip = (slice(None),) * axis
i = 0
for a in tup:
