code: 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, conv 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)
          )
          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([conv(1)], dtype=sdtype),
              ),
          )
  
      # Testing with argmin
      for sdtype, conv 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)
          )
          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([conv(1)], dtype=sdtype),
              ),
          )
  
imports: import pytest
  import tvm
  from tvm import IRModule, parser, relay, te
  from tvm.relay import analysis, op, transform
  from tvm.relay.op import op
  import numpy
label_desc: Use image node_id 33 for calling a global function with example usage: test_argreduce_infer_return_type() without return types