Could we specify or update output type?

I find we set handle type be float32, however, for some ops, like ArgMax, we need to make output’s dtype is int32. Can we have one mechanism to set it?
otherwise we will get error here:

Expr cond = (TVMArrayGet(UInt(8), handle, intrinsic::kArrTypeCode) ==
** UIntImm::make(UInt(8), dtype.code())** &&
TVMArrayGet(UInt(8), handle, intrinsic::kArrTypeBits) ==
UIntImm::make(UInt(8), dtype.bits()) &&
TVMArrayGet(UInt(16), handle, intrinsic::kArrTypeLanes) ==
UIntImm::make(UInt(16), dtype.lanes()));
Because handle type is float32 when input is float32, we can not change it.
Thanks

Currently, I disable the check TVMArrayGet(UInt(8), handle, intrinsic::kArrTypeCode) ==
** UIntImm::make(UInt(8), dtype.code()). but keep size and lanes check. Everything works fine. But I can not still do the work like this: https://github.com/dmlc/tvm/blob/master/topi/tests/python/test_topi_reduce.py#L42