Function Differences with tf.arg_max

View Source On Gitee

tf.arg_max

tf.arg_max(input, dimension, output_type=tf.dtypes.int64, name=None)

For more information, see tf.arg_max.

mindspore.Tensor.argmax

mindspore.Tensor.argmax(axis=None)

For more information, see mindspore.Tensor.argmax.

Usage

Same function. Two interfaces of MindSpore and TensorFlow decide on which dimension to return the index of the maximum value through the parameters axis and dimension, respectively.

The difference is that in the default state, axis=None of MindSpore returns the global index of the maximum value; TensorFlow’s dimension returns the maximum index of dimension=0 by default when no value is passed in.

Code Example

import mindspore as ms

a = ms.Tensor([[1, 10, 166.32, 62.3], [1, -5, 2, 200]], ms.float32)
print(a.argmax())
print(a.argmax(axis=0))
print(a.argmax(axis=1))
# output:
# 7
# [0 0 0 1]
# [2 3]

import tensorflow as tf
tf.enable_eager_execution()

b = tf.constant([[1, 10, 166.32, 62.3], [1, -5, 2, 200]])
print(tf.argmax(b).numpy())
print(tf.argmax(b, dimension=0).numpy())
print(tf.argmax(b, dimension=1).numpy())
# output:
# [0 0 0 1]
# [0 0 0 1]
# [2 3]