tf.math.argmax

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

View aliases

Main aliases

tf.argmax

tf.math.argmax(
    input,
    axis=None,
    output_type=tf.dtypes.int64,
    name=None
)

Used in the notebooks

Used in the guide	Used in the tutorials
Multilayer perceptrons for digit recognition with Core APIs Migrate metrics and optimizers Distributed training with Core APIs and DTensor Introduction to Tensors Introduction to Variables	Transfer learning with YAMNet for environmental sound classification Transfer learning with TensorFlow Hub Custom training: walkthrough Simple audio recognition: Recognizing keywords Distributed training with DTensors

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, 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>

Args
input	A Tensor.
axis	An integer, the axis to reduce across. Default to 0.
output_type	An optional output dtype (tf.int32 or tf.int64). Defaults to tf.int64.
name	An optional name for the operation.

Returns
A Tensor of type output_type.