Enhancement with deprecated_argument_lookup for argmax

This fix makes some enhancement for argmax, using
deprecated_argument_lookup instread of customerized logic.

Signed-off-by: Yong Tang <yong.tang.github@outlook.com>