tf.math.argmin

Returns the index with the smallest value across axes