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.