better support for numpy functions with 'axis' args