Use scalar for axis or None in MaxAndArgmax.

this remove a deprecation warning in NumPy as ndarray won't be
accepted as int anymore.