couldn't this check be moved to the main function too? like the axis != None case