[BugFix] Use shape dtype on ArgReduce to determine return type

Fix ArgReduce automatic return type inference by forcing it to use the
datatype of the shape of the Tensor instead of the fixed Int32.

Including additional tests.

Author: everton1984

src/relay/op/tensor/reduce.cc

@@ -338,7 +338,7 @@ bool GenericReduceRel(const Array<Type>& types, int num_inputs, const Attrs& att
 
   // assign output type and shape
   auto oshape = ReduceShapeImpl(in_shape, param, reporter);
-  reporter->Assign(types[1], TensorType(oshape, DataType::Int(32)));
+  reporter->Assign(types[1], TensorType(oshape, data->shape[0].dtype()));
   return true;
 }
 /*!

tests/python/relay/test_type_infer.py

@@ -23,6 +23,8 @@
 from tvm.relay import analysis, op, transform
 from tvm.relay.op import op as _op
 
+import numpy as np
+
 
 def infer_mod(mod, annotate_spans=True):
     if annotate_spans:
@@ -544,6 +546,42 @@ def test_repeat_register():
         assert "Operator custom_log3 is registered before" in str(cm.execption)
 
 
+def test_argreduce_infer_return_type():
+    x_shape = (1, 1)
+    broadcast_shape = [1, 1]
+    shape_dtypes = [("int32", lambda x: np.int32(x)), ("int64", lambda x: np.int64(x))]
+
+    # Testing with argmax
+    for sdtype in shape_dtypes:
+        x = relay.var("data", relay.TensorType(x_shape, "float32"))
+        broadcast_to = relay.op.broadcast_to(x, relay.const(broadcast_shape, dtype=sdtype[0]))
+        argmax = relay.op.argmax(broadcast_to, axis=[1])
+
+        f = relay.Function([x], argmax)
+        assert_has_type(
+            f,
+            relay.FuncType(
+                [relay.TensorType(broadcast_shape, "float32")],
+                relay.TensorType([sdtype[1](1)], dtype=sdtype[0]),
+            ),
+        )
+
+    # Testing with argmin
+    for sdtype in shape_dtypes:
+        x = relay.var("data", relay.TensorType(x_shape, "float32"))
+        broadcast_to = relay.op.broadcast_to(x, relay.const(broadcast_shape, dtype=sdtype[0]))
+        argmin = relay.op.argmin(broadcast_to, axis=[1])
+
+        f = relay.Function([x], argmin)
+        assert_has_type(
+            f,
+            relay.FuncType(
+                [relay.TensorType(broadcast_shape, "float32")],
+                relay.TensorType([sdtype[1](1)], dtype=sdtype[0]),
+            ),
+        )
+
+
 if __name__ == "__main__":
     import sys