relax
relax copied to clipboard
[Translator] Support translating op calls with Tuple input
Previously, when a Relay function contains a Call which directly uses Tuples as arguments (the example below),
%25 = (%23, %24) /* ty=(Tensor[(1, 160), float32], Tensor[(1, 160), float32]) */;
%26 = concatenate(%25, axis=-1) /* ty=Tensor[(1, 320), float32] */;
our Relay-translator is unable to generate corresponding CallTIR, because the translator always assumes a argument of a Call is mapped to a single tensor (see the code snippet below: the translator directly passes the Relax variable new_args[-1]
to function te_tensors
, which translate a Var to a single tensor).
https://github.com/tlc-pack/relax/blob/60e9a01cdfdd013945790fc03d5abad29b8a7c0b/python/tvm/relax/testing/relay_translator.py#L124
https://github.com/tlc-pack/relax/blob/60e9a01cdfdd013945790fc03d5abad29b8a7c0b/src/relax/ir/emit_te.h#L56-L61
But in fact, the Relax variable may correspond to a Tuple of tensors, which wasn’t taken into consideration before. And such case can lead to error in TETensor
, when creating tensors.
Therefore, this PR fixes the issue by examine the Relax variable before the tensor creation of Relay Call arguments. If an argument has shape Tuple and type TupleType, we break down the tuple Variable and emit a TupleGetItem for each field, and meanwhile create a tensor for each field.
cc @YuchenJin
cc @MasterJH5574 would be great to rebase and get it in