revert changes to cupy.linalg.norm
changes to cub reductions made setting keepdims=False and axis=None unecessary
@@ -61,20 +61,18 @@ def norm(x, ord=None, axis=None, keepdims=False):
axis = (axis,)
if len(axis) == 1:
- if nd == 1:
- # fast CUB-based reductions require axis=None, keepdims=False
- axis = None
if ord == numpy.Inf:
- ret = abs(x).max(axis=axis)
+ return abs(x).max(axis=axis, keepdims=keepdims)
elif ord == -numpy.Inf:
- ret = abs(x).min(axis=axis)
+ return abs(x).min(axis=axis, keepdims=keepdims)
elif ord == 0:
# Zero norm
# Convert to Python float in accordance with NumPy
- ret = (x != 0).astype(x.real.dtype).sum(axis=axis)
+ return (x != 0).astype(x.real.dtype).sum(
+ axis=axis, keepdims=keepdims)
elif ord == 1:
# special case for speedup
- ret = abs(x).sum(axis=axis)
+ return abs(x).sum(axis=axis, keepdims=keepdims)
elif ord is None or ord == 2:
# special case for speedup
if x.dtype.kind == 'c':
@@ -82,7 +80,7 @@ def norm(x, ord=None, axis=None, keepdims=False):
s *= s
else:
s = x * x
- ret = cupy.sqrt(s.sum(axis=axis))
+ return cupy.sqrt(s.sum(axis=axis, keepdims=keepdims))
else:
try:
float(ord)
@@ -91,14 +89,8 @@ def norm(x, ord=None, axis=None, keepdims=False):
absx = abs(x)
absx **= ord
- ret = absx.sum(axis=axis)
+ ret = absx.sum(axis=axis, keepdims=keepdims)
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
