def normalize(v, axis=-1, order=2):
"""
Normalize a vector for arbitrary axes, and giving optimal performance.
"""
one_dim = np.ndim(v) == 0 or np.ndim(v) == 1
l2 = np.atleast_1d(np.linalg.norm(v, order, axis))
l2[l2==0] = 1
ret = a / np.expand_dims(l2, axis)
return ret[0] if one_dim else ret
This code will work: