tf.math.argmax

Returns the index with the largest value across axes of a tensor.

Used in the notebooks

Used in the guide Used in the tutorials

In case of identity returns the smallest index.

For example:

A = tf.constant([2, 20, 30, 3, 6])
tf.math.argmax(A)  # A[2] is maximum in tensor A
<tf.Tensor: shape=(), dtype=int64, numpy=2>
B = tf.constant([[2, 20, 30, 3, 6], [3, 11, 16, 1, 8],
                 [14, 45,