* Adapt tensor axis for argmin/max

* Add UT

* Polish UT

* move arg min max to phi.

* move infermeta.

* fix as reviews.