If the 'axis' variable is a numpy scalar array (like numpy.array(1)), these conditional branches will fail as it tries to build a set over the non-iterable array. I've added another branch to handle this case.