tf.argmax(input,axis)根据axis取值的不同返回每行或者每列最大值的索引。
代码如下：
import tensorflow as tf
import numpy as np
sess=tf.Session()
a = np.array([[1, 2, 3], [2, 3, 4], [5, 4, 3], [8, 7, 2]])
a0=tf.argmax(a,axis=0)
a1=tf.argmax(a,axis=1)
a0=sess.run(a0)
a1=sess.run(a1)
b = np.array([[[1, 2, 3,5], [2, 3, 4,8], [5, 1,4, 3]],
[[13 ,4, 3,5], [2,13, 4,8],[8, 4, 7, 32]]
])
b0=tf.argmax(b,axis=0)
b1=tf.argmax(b,axis=1)
b2=tf.argmax(b,axis=2)
b0=sess.run(b0)
b1=sess.run(b1)
b2=sess.run(b2)
print('矩阵为二维度状况：\n')
print('first dimension=',a0)
print('second dimension=',a1)
print('矩阵为三个维度状况：\n')
print('first dimension=',b0)
print('second dimension=',b1)
print('third dimension=',b2)
结果如下：
