tf.math.argmax  |  TensorFlow v2.16.1 (original) (raw)

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

In case of identity returns the smallest index.

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, 23, 5, 27]]) tf.math.argmax(B, 0) <tf.Tensor: shape=(5,), dtype=int64, numpy=array([2, 2, 0, 2, 2])> tf.math.argmax(B, 1) <tf.Tensor: shape=(3,), dtype=int64, numpy=array([2, 2, 1])> C = tf.constant([0, 0, 0, 0]) tf.math.argmax(C) # Returns smallest index in case of ties <tf.Tensor: shape=(), dtype=int64, numpy=0>