Created by: baojun-nervana

Softmax op was updated to add the axis arg. this PR is to update and reflect the change