TensorFlow: argmax (-min)

Question

I just noticed an unexpected (at least for me) behavior in TensorFlow. I thought tf.argmax (-argmin) operates on the ranks of a Tensor from outer to inner, but apparently it does not?!

Example:

import numpy as np
import tensorflow as tf

sess = tf.InteractiveSession()

arr = np.array([[31, 23,  4, 24, 27, 34],
                [18,  3, 25,  0,  6, 35],
                [28, 14, 33, 22, 20,  8],
                [13, 30, 21, 19,  7,  9],
                [16,  1, 26, 32,  2, 29],
                [17, 12,  5, 11, 10, 15]])

# arr has rank 2 and shape (6, 6)
tf.rank(arr).eval()
> 2
tf.shape(arr).eval()
> array([6, 6], dtype=int32)

tf.argmax takes two arguments: input and dimension. Since the indices of array arr are arr[rows, columns], I would expect tf.argmax(arr, 0) to return the index of the maximum element per row, while I would have expected tf.argmax(arr, 1) to return the maximum element per column. Likewise for tf.argmin.

However, the opposite is true:

tf.argmax(arr, 0).eval()
> array([0, 3, 2, 4, 0, 1])

# 0 -> 31 (arr[0, 0])
# 3 -> 30 (arr[3, 1])
# 2 -> 33 (arr[2, 2])
# ...
# thus, this is clearly searching for the maximum element
# for every column, and *not* for every row

tf.argmax(arr, 1).eval()
> array([5, 5, 2, 1, 3, 0])

# 5 -> 34 (arr[0, 5])
# 5 -> 35 (arr[1, 5])
# 2 -> 33 (arr[2, 2])
# ...
# this clearly returns the maximum element per row,
# albeit 'dimension' was set to 1

Can someone explain this behavior?

Generalized every n-dimensional Tensor t is indexed by t[i, j, k, ...]. Thus, t has rank n and shape (i, j, k, ...). Since dimension 0 corresponds to i, dimension 1 to j, and so forth. Why does tf.argmax (& -argmin) ignore this scheme?

Answer 1

Think of the dimension argument of tf.argmax as the axis across which you reduce. tf.argmax(arr, 0) reduces across dimension 0, i.e. the rows. Reducing across rows means that you will get the argmax of each individual column.

This might be counterintuitive, but it falls in line with the conventions used in tf.reduce_max and so on.