[Bug] Compilation failure for broadcast-argmax in internal type inference

[ganler]

Expected behavior

Pass compilation.

Actual behavior

Having:

"""
fn (%i0: Tensor[(2, 1), float64]) {
  %0 = broadcast_to(%i0, shape=[2i64, 1i64]);
  %1 = argmax(%0, axis=[0]);
  (%0, %1)
}
"""

will fail type inference used in some intermediate passes. But it is a well-formed model and can pass a direct type inference. So some passes might create ill-formed relay IR and reveals the error in some later type inference.

Environment

fa17da2 on Ubuntu 20.04.

Steps to reproduce

Preferably a minimal script to cause the issue to occur.

import tvm
from tvm import relay

"""
fn (%i0: Tensor[(2, 1), float64]) {
  %0 = broadcast_to(%i0, shape=[2i64, 1i64]);
  %1 = argmax(%0, axis=[0]);
  (%0, %1)
}
"""

i0 = relay.var("i0", shape=[2, 1], dtype="float64")
v0 = relay.broadcast_to(i0, shape=relay.const([2, 1], dtype="int64"))
v1 = relay.argmax(v0, axis=[0])
v2 = relay.Tuple([v0, v1])

f = relay.Function([i0], v2)
relay.create_executor("graph", device=tvm.cpu(), target="llvm").evaluate(f)

"""
The Relay type checker is unable to show the following types match.
In particular `Tensor[(1i64), int64]` does not match `Tensor[(1i64), int32]`
"""

Triage

Please refer to the list of label tags linked above to find the relevant tags and add them here in a bullet format (example below).

• needs-triage