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`.

© 2022 The TensorFlow Authors. All rights reserved.
Licensed under the Creative Commons Attribution License 4.0.
Code samples licensed under the Apache 2.0 License.
https://www.tensorflow.org/api_docs/python/tf/math/argmax