Updated check logic
@@ -44,14 +44,11 @@ def _normalize_axis_index(axis, ndim):
The normalized axis index, such that `0 <= normalized_axis < ndim`
"""
- if axis < 0 and abs(axis) > ndim:
+ if ndim <= axis or (axis < -ndim and axis + ndim < 0):
raise _errors._AxisError('axis '+str(axis)+' is out of bounds for ' +
'array of dimension '+str(ndim))
- elif axis > 0 and abs(axis) >= ndim:
- raise _errors._AxisError('axis '+str(axis)+' is out of bounds for ' +
- 'array of dimension '+str(ndim))
- else:
+ elif (0 <= axis < ndim) or (-ndim <= axis < 0 and 0 <= axis + ndim < ndim):
if axis < 0:
return (ndim - abs(axis)) % ndim
else:
