BUG: Made use of np.apply_along_axis and replaced custom version