Skip to content

Conversation

@PatrikPerssonInceptron
Copy link
Contributor

@PatrikPerssonInceptron PatrikPerssonInceptron commented Feb 11, 2025

Replaces topi.split with relax.op.split in the ONNX frontend to align with the Relax operator API.

Fixes related ONNX frontend unit tests by:

  • Updating InferStructInfoSplit to return a TensorStructInfo instead of a TupleStructInfo if the operator produces only one split, to be consistent with topi.split
  • Updating the relax.split legalize callback to only perform the modulo != 0 check when the modulo result is a constant expression

Fixes split unit tests when only one tensor is returned by changing the expected struct info from a TupleStructInfo to a TensorStructInfo.

The reason why it is necessary to update InferStructInfoSplit to return a TensorStructInfo instead of a TupleStructInfo when the split operator only returns a single tensor is because the MakeCallTIR api will unpack any tuple of struct info containing only one element. See tvm/src/relax/op/op.cc line 577:

Expr MakeCallTIR(Expr func, Tuple args, Array<TensorStructInfo> out_sinfo_list,
                 Optional<Expr> packed_ints) {
  ...
  StructInfo out_sinfo{nullptr};
  if (out_sinfo_list.size() == 1) {
    out_sinfo = out_sinfo_list[0];
  } else {
    out_sinfo = TupleStructInfo({out_sinfo_list.begin(), out_sinfo_list.end()});
  }

fixed related onnx frontend unit tests
@Hzfengsy Hzfengsy merged commit f140fb4 into apache:main Feb 16, 2025
14 checks passed
@Hzfengsy
Copy link
Member

Thanks @PatrikPerssonInceptron

ShiboXing pushed a commit to ShiboXing/tvm that referenced this pull request Aug 10, 2025
…tend (apache#17642)

* replaced topi.split with relax.op.split in onnx frontend

fixed related onnx frontend unit tests

* updated formatting

* updated the split unit tests for the case when only one tensor is returned

* formatting
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants