TF argmax - handling int64 datatype (#6674)

Co-authored-by: Ubuntu <ubuntu@ip-172-31-0-202.us-west-2.compute.internal>

Author: anijain2305

python/tvm/relay/frontend/tensorflow.py

@@ -146,7 +146,11 @@ def _impl(inputs, attr, params, mod):
             raise TypeError(
                 "Unsupported argument for `{}` : `axis` should be a constant".format(func_name)
             )
-        return func(inputs[0], axis=axis_input_value, keepdims=False)
+        out = func(inputs[0], axis=axis_input_value, keepdims=False)
+        dtype = attr["output_type"].name
+        if dtype != "int32":
+            out = _op.cast(out, dtype=dtype)
+        return out
 
     return _impl
 

tests/python/frontend/tensorflow/test_forward.py

@@ -1601,16 +1601,16 @@ def _test_argx(func, data, **kwargs):
 
     with tf.Graph().as_default():
         inp = array_ops.placeholder(shape=data.shape, dtype=data.dtype, name="c0")
-        func(inp, name="argx0", output_type=tf.int32, **kwargs)
-
+        func(inp, name="argx0", **kwargs)
         compare_tf_with_tvm(data, "c0:0", "argx0:0")
 
 
 def test_forward_argminmax():
-    for axis in [None, 0, 1, 2]:
-        data = np.random.uniform(size=(8, 4, 9)).astype("float32")
-        _test_argx(tf.argmax, data=data, axis=axis)
-        _test_argx(tf.argmin, data=data, axis=axis)
+    for output_type in [tf.int64, tf.int32]:
+        for axis in [None, 0, 1, 2]:
+            data = np.random.uniform(size=(8, 4, 9)).astype("float32")
+            _test_argx(tf.argmax, data=data, axis=axis, output_type=output_type)
+            _test_argx(tf.argmin, data=data, axis=axis, output_type=output_type)
 
 
 #######################################################################