[Fix][Relax][Frontend][PyTorch] Fix x.split(int) with a non-divisible split_size

[siyiweigeHEW]

[Relax][Frontend][PyTorch] Fix x.split(int) with a non-divisible split_size

Fixes: #20232

Summary

torch.split(x, s, dim) splits dim into chunks of size s, with the
last chunk smaller when the dimension D is not divisible by s. The
Relax PyTorch frontend's _split converter
(base_fx_graph_translator.py) converted the per-chunk size into a
section count n_section = ceil(D / s) and passed it to
relax.op.split's integer argument — whose semantics are "split into
n_section equal sections" (split_len = ceil(D / n_section),
src/relax/op/tensor/manipulate.cc). Whenever
ceil(D / ceil(D / s)) != s (e.g. split_size > D/2 with a
non-divisible D), valid PyTorch models silently produced
differently-shaped chunks. For example x.split(6) on a (10,) tensor
yielded (5,), (5,) instead of (6,), (4,).

This PR converts the int per-chunk size into the cumulative cut
positions [s, 2s, ..., (ceil(D/s) - 1) * s] — the same indices form
the list/tuple branch already passes to relax.op.split — so both
forms produce PyTorch-identical chunk shapes.

Root cause

_split handles two aten ops: split.Tensor (int split_size) and
split_with_sizes.default (list/tuple). The list branch builds
cumulative cut positions and is correct. The int branch instead computed
n_section = ceil(D / split_size) and relied on relax.op.split's
integer "equal sections" semantics, which only coincide with PyTorch's
per-chunk-size semantics when ceil(D / ceil(D / s)) == s (divisible
sizes, or e.g. D=10, s=3). The bug is a semantic mismatch between
"chunks of size s" (PyTorch) and "ceil(D/s) equal sections"
(relax.op.split int), not a numerical issue.

Fix

python/tvm/relax/frontend/torch/base_fx_graph_translator.py —
_split, int branch:

else:
    # torch.split(x, s, dim) splits dim into chunks of size s, with the
    # last chunk smaller if D % s != 0. relax.op.split's integer argument
    # is the number of *equal* sections, so passing ceil(D / s) yields
    # wrong shapes whenever ceil(D / ceil(D / s)) != s (e.g. s > D/2).
    # Convert the per-chunk size to the cumulative cut positions instead,
    # mirroring the list/tuple branch above.
    dim_size = self.shape_of(x)[dim].value
    num_chunks = (dim_size + split_size - 1) // split_size
    n_section = [split_size * i for i in range(1, num_chunks)]

The list/tuple branch and the split_with_sizes.default mapping are
unchanged. _split is shared by from_exported_program and from_fx
(via BaseFXGraphImporter), so both entry points are covered.

Validation

In-tree regression test (added)

test_split_int_split_size in
tests/python/relax/test_frontend_from_exported_program.py:

• structural check: x.split(6, dim=0) on a (10,) input lowers to
  R.split(input, indices_or_sections=[6], axis=0) with output shapes
  (6,), (4,) (asserted via verify_model structural equality);
• numerical check vs native PyTorch over non-divisible sizes and dims:
  (10,) s=6/7/8/9 dim=0, (12,) s=7 dim=0, (12,8) s=5 dim=1,
  (3,10) s=6 dim=-1 — shapes and values all match.

Differential test

verify_patch.py (in prove_hum/torch_split/) runs the full suite on
the fixed _split injected verbatim from this branch against the
pre-fix _split from origin/main (tvm-env 0.18, whose
_split is byte-identical to the current frontend):

• before: 10/10 diverging cases (non-divisible s > D/2 across
  dim=0, dim=1, and negative dims) reproduce the bug, e.g.
  (10,) s=6: torch [(6,), (4,)] vs TVM [(5,), (5,)];
• after: all 10 now match torch shapes and values; the divisible
  baseline (4/4) and the list/tuple control group (3/3) remain
  unchanged — no regression.

Known pre-existing limitation (unchanged by this PR, same family as the
ONNX single-output Split handling): split_size >= D or a
single-element split_with_sizes produces a single chunk, which older
relax versions cannot import as a 1-tuple (single-output struct-info
inference). On current main the empty-indices form already lowers to a
proper 1-tuple, and this PR keeps int split_size >= D consistent with
the existing single-element list behavior.

Files changed

• python/tvm/relax/frontend/torch/base_fx_graph_translator.py — fix
  the int split_size branch of _split to pass cumulative cut
  positions instead of a section count.
• tests/python/relax/test_frontend_from_exported_program.py — add
  test_split_int_split_size regression coverage.
• tests/python/relax/test_frontend_from_fx.py — test_split's
  expected1 asserted the old int-section IR form
  (indices_or_sections=3 for x.split(1, dim=1)); updated to the
  correct list form (indices_or_sections=[1, 2]).