Question about numpy.apply_along_axis