[ONNX] Fix Split operator for uneven splits with num_outputs

[Mr-Neutr0n]

Summary

Fixes #18751

The ONNX frontend's Split converter crashes when the tensor dimension isn't evenly divisible by num_outputs. For example, splitting a length-10 tensor into 3 parts should produce [4, 4, 2] per the ONNX Opset 18 spec, but the converter only handled perfectly divisible cases.

Root cause: When no explicit split sizes are provided, both _impl_v1 and _impl_v13 pass the raw num_outputs integer to relax.op.split. While the underlying op supports uneven splits in theory, the conversion fails during model normalization for non-divisible dimensions.

Fix: Added a _compute_split_indices helper that computes explicit cumulative split indices using ceil-based block sizes when the input shape is statically known:

block_size = ceil(dim / num_outputs)
indices = [block_size, 2*block_size, ..., (N-1)*block_size]

For a length-10 tensor with num_outputs=3: block_size = ceil(10/3) = 4, indices = [4, 8], producing splits of [4, 4, 2]. Falls back to passing the integer for dynamic shapes.

Test plan

• Added test cases for 1D uneven split (10 / 3 -> [4, 4, 2])
• Added test cases for 2D uneven split along axis 1 (7 / 3 -> [3, 3, 1])
• Existing even-split tests remain unchanged and should continue to pass

[gemini-code-assist]

Summary of Changes

Hello @Mr-Neutr0n, I'm Gemini Code Assist¹! I'm currently reviewing this pull request and will post my feedback shortly. In the meantime, here's a summary to help you and other reviewers quickly get up to speed!

This pull request addresses a critical bug in the ONNX frontend's Split operator within TVM Relax, which previously failed when encountering uneven tensor dimensions combined with the num_outputs attribute. The core of the solution involves dynamically computing precise split indices based on the ONNX Opset 18 specification, allowing for correct handling of non-divisible tensor splits. This enhancement significantly improves the robustness and compliance of the ONNX frontend, preventing crashes and ensuring accurate tensor partitioning in more complex scenarios.

Highlights

• ONNX Split Operator Fix: Resolved an issue where the ONNX frontend's Split operator crashed when attempting to split a tensor into an uneven number of parts using num_outputs, which previously only supported perfectly divisible cases.
• Uneven Split Handling: Implemented a new helper method, _compute_split_indices, to calculate explicit cumulative split indices for static input shapes, adhering to the ONNX Opset 18 specification for ceil-based block sizes (e.g., splitting a length-10 tensor into 3 parts yields [4, 4, 2]).
• Test Coverage: Added new test cases to validate the correct behavior of the Split operator for both 1D and 2D uneven splits, ensuring compliance with the ONNX Opset 18 specification.

🧠 New Feature in Public Preview: You can now enable Memory to help Gemini Code Assist learn from your team's feedback. This makes future code reviews more consistent and personalized to your project's style. Click here to enable Memory in your admin console.

Changelog

• python/tvm/relax/frontend/onnx/onnx_frontend.py
  • Added _compute_split_indices method to correctly calculate split indices for uneven splits based on ONNX Opset 18.
  • Updated _impl_v1 and _impl_v13 to utilize the new _compute_split_indices method when split attribute is not provided.
• tests/python/relax/test_frontend_onnx.py
  • Added new test cases for 1D and 2D uneven splits to verify the fix.

Activity

• No activity found for this pull request.

Using Gemini Code Assist

The full guide for Gemini Code Assist can be found on our documentation page, here are some quick tips.

Invoking Gemini

You can request assistance from Gemini at any point by creating a comment using either /gemini <command> or @gemini-code-assist <command>. Below is a summary of the supported commands on the current page.

Feature	Command	Description
Code Review	/gemini review	Performs a code review for the current pull request in its current state.
Pull Request Summary	/gemini summary	Provides a summary of the current pull request in its current state.
Comment	@gemini-code-assist	Responds in comments when explicitly tagged, both in pull request comments and review comments.
Help	/gemini help	Displays a list of available commands.

Customization

To customize Gemini Code Assist for GitHub experience, repository maintainers can create a configuration file and/or provide a custom code review style guide (such as PEP-8 for Python) by creating and adding files to a .gemini/ folder in the base of the repository. Detailed instructions can be found here.

Limitations & Feedback

Gemini Code Assist may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counter productive. You can react with 👍 and 👎 on @gemini-code-assist comments. If you're interested in giving your feedback about your experience with Gemini Code Assist for Github and other Google products, sign up here.

You can also get AI-powered code generation, chat, as well as code reviews directly in the IDE at no cost with the Gemini Code Assist IDE Extension.

Footnotes

1. Review the Privacy Notices, Generative AI Prohibited Use Policy, Terms of Service, and learn how to configure Gemini Code Assist in GitHub here. Gemini can make mistakes, so double check it and use code with caution. ↩

[gemini-code-assist]

Code Review

This pull request correctly fixes a bug in the ONNX Split operator converter for uneven splits when num_outputs is provided. The introduction of the _compute_split_indices helper function to handle static shapes aligns with the ONNX specification and is a clean solution. The new test cases for 1D and 2D uneven splits are also a great addition, ensuring the fix is robust. I have a couple of minor suggestions to improve code readability by avoiding redundant dictionary lookups.

[gemini-code-assist]

To improve readability and avoid a redundant dictionary lookup, you can reuse the axis variable that was already retrieved on line 1824.

Suggested change

	return relax.op.split(inputs[0], indices, attr.get("axis", 0))
	return relax.op.split(inputs[0], indices, axis)

[gemini-code-assist]

Similar to the change in _impl_v1, you can reuse the axis variable that was retrieved on line 1847 to improve readability and avoid a redundant dictionary lookup.

Suggested change

	return relax.op.split(inputs[0], indices, attr.get("axis", 0))
	return relax.op.split(inputs[0], indices, axis)

[locnd182644]

@Mr-Neutr0n Thank you about this PR. But it seems the "uneven splits" issue has been solved in PR #17757 and apply at TVM Version: 0.20.0. You can check again.

[tlopex]

Hi @Mr-Neutr0n Could you look at the error in the CI and fix it?

[Mr-Neutr0n]

Two notes:

1. Pushed c0061069 — small refactor reusing the axis local in _impl_v1/_impl_v13 per @gemini-code-assist's suggestion. Pure cleanup, no behavior change.

2. On the cpu/pr-head build error — I have to be honest: I can't run TVM's C++ build locally to debug it (the full build is multi-hour on a workstation and would burn the rest of the day). The error log itself isn't visible from the GitHub side; ci.tlcpack.ai returns a redirect without the failing test name in the API response. The Python change itself is straightforward (one helper function, no C++ surface), so the build failure is most likely a stale build cache or an unrelated flake — but I don't want to guess and burn another CI cycle.

@locnd182644 — you said back in February that "the 'uneven splits' issue has been solved in PR #17757 and apply at TVM Version: 0.20.0. You can check again." The original issue (#18751) is still open, though, so I want to confirm: is #17757 actually sufficient to close #18751, in which case this PR is redundant and I'd rather close it cleanly, or is the ONNX frontend path genuinely a separate gap? Happy either way — just let me know which.

[Mr-Neutr0n]

Rebased and resolved the conflicts.

[Mr-Neutr0n]

Correction first: my comment yesterday said "Rebased and resolved the conflicts." That was wrong — the push never landed, and the branch sat at c0061069, 640 commits behind main and still conflicting. Sorry for the noise; it's actually done now and the PR shows MERGEABLE.

While redoing it properly I found a real bug in my own refactor commit, so this is worth more than a rebase note.

c0061069 was broken

That commit applied @gemini-code-assist's suggestion to reuse the axis local, and I got it wrong in both _impl_v1 and _impl_v13:

        else:
            axis = attr.get("axis", 0)      # only bound HERE
            ...
        return relax.op.split(inputs[0], indices, axis)   # used unconditionally

axis is assigned only inside the else, so any Split node with an explicit split list — the common path, and every pass_split=True test case — would raise UnboundLocalError. It slipped through because the cpu/pr-head build was already failing for unrelated reasons, so nothing ever exercised it.

Fixed by hoisting axis = attr.get("axis", 0) to the top of both methods. Good argument for not applying a suggested refactor without a green run behind it.

The fix is now narrower

Original version returned explicit indices for every static num_outputs split, so even splits changed from indices_or_sections=3 to [2, 4] — semantically identical, but it churned the IR (and the expectations) for cases that already worked. It now keeps the integer form unless the split is genuinely uneven:

if dim % num_outputs:
    block_size = math.ceil(dim / num_outputs)
    return [block_size * i for i in range(1, num_outputs)]
return num_outputs

So the emitted IR is byte-identical to main for everything except the uneven case this PR exists to fix.

Test conflict

test_frontend_onnx.py was restructured on main since February — verify_split now takes an expected IRModule from a make_expected helper, driven by a split_cases table. I took main's version wholesale and re-expressed my two cases as table entries:

# 10 / 3 -> [4, 4, 2]
(10, [[4], [4], [2]], False, 0, False, 18),
# 7 / 3 along axis 1 -> [3, 3, 1]
((4, 7), [[4, 3], [4, 3], [4, 1]], False, 1, False, 18),

make_expected needed teaching about this, since its pass_split=False branch always produced indices_or_sections=len(outdata_shapes), which is exactly the wrong expectation for an uneven split. It now mirrors the frontend: explicit indices when the axis is static and uneven, the integer otherwise.

What I could and couldn't verify

Being upfront, as before: I still cannot build TVM locally, so I have not run the suite. What I did instead was re-implement both sides of the contract in plain Python and check they agree across every entry in split_cases, both dynamic values:

shape=6      axis=0 dynamic=False  frontend=3       expected=3
shape=3      axis=0 dynamic=False  frontend=3       expected=3
shape=10     axis=0 dynamic=False  frontend=[4, 8]  expected=[4, 8]
shape=(4, 7) axis=1 dynamic=False  frontend=[3, 6]  expected=[3, 6]
...all agree: True

and that the declared output shapes match a ceil-based split (10 -> [4,4,2], 7 -> [3,3,1]). That is a check of the arithmetic, not of TVM — CI remains the real verdict.

Known limitation

For a dynamic axis the converter still falls back to the integer form, so an uneven split with an unknown dimension remains incorrect. That can't be lowered to a static index list; it would need shape-dependent lowering. Out of scope here, and unchanged from main — flagging it so the fix isn't read as more complete than it is.

[gemini-code-assist]

Caution

The consumer version of Gemini Code Assist on GitHub has been sunset. All code review activity has officially ceased.