In [ ]:
# Created 2016-04-10
# Tensorflow version: 0.7

In [2]:
import numpy as np
import tensorflow as tf

In [3]:
# Although this is for reduce_max and argmax, the same goes 
# for reduce_min and argmin.
#
# There is a trick to correctly find the dimension to be specified
# in all of those "reduce" like functions:
# It is the dimension that will disappear.

a = tf.constant([[1, 2, 6], [0, 8, 5]])

# The max along the 0th dimension is [1, 8, 6].
# The corresponding indices in the 1st dimension can be obtained by
# argmax, which results to [0 1 0]
max_along_0 = tf.reduce_max(a, 0)
argmax_along_0 = tf.argmax(a, 0)

# The max along the 1st dimension is [6, 8].
# The corresponding indices in the 0th dimension can be obtained by
# argmax, which results to [2 1]
max_along_1 = tf.reduce_max(a, 1)
argmax_along_1 = tf.argmax(a, 1)


with tf.Session() as sess:
    print('--------- dimension 0 ----------')
    print(sess.run(max_along_0))
    print(sess.run(argmax_along_0))
    
    print('--------- dimension 1 ----------')
    print(sess.run(max_along_1))
    print(sess.run(argmax_along_1))


--------- dimension 0 ----------
[1 8 6]
[0 1 0]
--------- dimension 1 ----------
[6 8]
[2 1]