enable fast CUB-based 1D path for keepdims=True as well
@@ -62,19 +62,19 @@ def norm(x, ord=None, axis=None, keepdims=False):
if len(axis) == 1:
if nd == 1:
- axis = None # fast CUB-based reductions require axis == None
+ # fast CUB-based reductions require axis=None, keepdims=False
+ axis = None
if ord == numpy.Inf:
- return abs(x).max(axis=axis, keepdims=keepdims)
+ ret = abs(x).max(axis=axis)
elif ord == -numpy.Inf:
- return abs(x).min(axis=axis, keepdims=keepdims)
+ ret = abs(x).min(axis=axis)
elif ord == 0:
# Zero norm
# Convert to Python float in accordance with NumPy
- return (x != 0).astype(x.real.dtype).sum(
- axis=axis, keepdims=keepdims)
+ ret = (x != 0).astype(x.real.dtype).sum(axis=axis)
elif ord == 1:
# special case for speedup
- return abs(x).sum(axis=axis, keepdims=keepdims)
+ ret = abs(x).sum(axis=axis)
elif ord is None or ord == 2:
# special case for speedup
if x.dtype.kind == 'c':
@@ -82,7 +82,7 @@ def norm(x, ord=None, axis=None, keepdims=False):
s *= s
else:
s = x * x
- return cupy.sqrt(s.sum(axis=axis, keepdims=keepdims))
+ ret = cupy.sqrt(s.sum(axis=axis))
else:
try:
float(ord)
@@ -91,8 +91,14 @@ def norm(x, ord=None, axis=None, keepdims=False):
absx = abs(x)
absx **= ord
- ret = absx.sum(axis=axis, keepdims=keepdims)
+ ret = absx.sum(axis=axis)
ret **= cupy.reciprocal(ord, dtype=ret.dtype)
+ if keepdims:
+ if nd == 1:
+ return ret.reshape((1,))
+ ret_shape = list(x.shape)
+ ret_shape[axis[0]] = 1
+ ret = ret.reshape(ret_shape)
return ret
elif len(axis) == 2:
row_axis, col_axis = axis
