Optimizer CPU code
@@ -51,6 +51,8 @@ def concatenate(tup, axis=0):
"""
ndim = None
shape = None
+ dtype = None
+ have_same_types = True
for a in tup:
if not isinstance(a, cupy.ndarray):
raise TypeError('Only cupy arrays can be concatenated')
@@ -60,13 +62,15 @@ def concatenate(tup, axis=0):
ndim = a.ndim
shape = list(a.shape)
axis = _get_positive_axis(a.ndim, axis)
+ dtype = a.dtype
continue
+ have_same_types &= (a.dtype == dtype)
if a.ndim != ndim:
raise ValueError(
'All arrays to concatenate must have the same ndim')
- if any(i != axis and shape[i] != a.shape[i]
- for i in six.moves.range(ndim)):
+ for i in six.moves.range(ndim):
+ if i != axis and shape[i] != a.shape[i]:
raise ValueError(
'All arrays must have same shape except the axis to '
'concatenate')
@@ -75,6 +79,7 @@ def concatenate(tup, axis=0):
if ndim is None:
raise ValueError('Cannot concatenate from empty tuple')
+ if not have_same_types:
dtype = numpy.find_common_type([a.dtype for a in tup], [])
return core.concatenate(tup, axis, shape, dtype)
