This is an automated email from the ASF dual-hosted git repository.
tlopex pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/tvm.git
The following commit(s) were added to refs/heads/main by this push:
new 2dfd3cc200 [Fix][Relax][Frontend][PyTorch] Fix `x.split(int)` with a
non-divisible `split_size` (#20240)
2dfd3cc200 is described below
commit 2dfd3cc200ace10e56264abcb9aa3ecd048849f8
Author: HuEnwei <[email protected]>
AuthorDate: Wed Sep 2 02:55:32 2026 +0800
[Fix][Relax][Frontend][PyTorch] Fix `x.split(int)` with a non-divisible
`split_size` (#20240)
# [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:
```python
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]`).
---
.../frontend/torch/base_fx_graph_translator.py | 10 ++-
.../relax/test_frontend_from_exported_program.py | 80 ++++++++++++++++++++++
tests/python/relax/test_frontend_from_fx.py | 2 +-
3 files changed, 90 insertions(+), 2 deletions(-)
diff --git a/python/tvm/relax/frontend/torch/base_fx_graph_translator.py
b/python/tvm/relax/frontend/torch/base_fx_graph_translator.py
index d600987cdd..2afbd009e7 100644
--- a/python/tvm/relax/frontend/torch/base_fx_graph_translator.py
+++ b/python/tvm/relax/frontend/torch/base_fx_graph_translator.py
@@ -2330,7 +2330,15 @@ class BaseFXGraphImporter(metaclass=abc.ABCMeta):
cum_sum = 0 if not n_section else n_section[-1]
n_section.append(s + cum_sum)
else:
- n_section = (self.shape_of(x)[dim].value + split_size - 1) //
split_size
+ # 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)]
return self.block_builder.emit(relax.op.split(x, n_section, dim))
def _squeeze(self, node: fx.Node) -> relax.Var:
diff --git a/tests/python/relax/test_frontend_from_exported_program.py
b/tests/python/relax/test_frontend_from_exported_program.py
index 7dc3c73564..53e8a3a314 100644
--- a/tests/python/relax/test_frontend_from_exported_program.py
+++ b/tests/python/relax/test_frontend_from_exported_program.py
@@ -5842,6 +5842,86 @@ def test_split():
verify_model(Chunk(), example_args, {}, Expected)
+def test_split_int_split_size():
+ """x.split(int, dim) must produce chunks of size `split_size` (the last one
+ smaller when the dimension is not divisible), matching PyTorch.
+
+ The frontend used to convert the int per-chunk size into a section count
and
+ pass it as relax.op.split's int argument, which means "split into N equal
+ sections"; that yields wrong chunk shapes whenever
+ ceil(D / ceil(D / split_size)) != split_size (e.g. split_size > D/2). The
+ int branch now builds cumulative cut positions, the same as the list/tuple
+ form.
+ """
+
+ class Split6(Module):
+ def forward(self, input):
+ return input.split(6, dim=0)
+
+ @tvm.script.ir_module
+ class Expected:
+ @R.function
+ def main(input: R.Tensor((10,), dtype="float32")) -> R.Tuple(
+ R.Tensor((6,), dtype="float32"),
+ R.Tensor((4,), dtype="float32"),
+ ):
+ with R.dataflow():
+ lv: R.Tuple(
+ R.Tensor((6,), dtype="float32"),
+ R.Tensor((4,), dtype="float32"),
+ ) = R.split(input, indices_or_sections=[6], axis=0)
+ lv1: R.Tensor((6,), dtype="float32") = lv[0]
+ lv2: R.Tensor((4,), dtype="float32") = lv[1]
+ gv: R.Tuple(
+ R.Tensor((6,), dtype="float32"),
+ R.Tensor((4,), dtype="float32"),
+ ) = (lv1, lv2)
+ R.output(gv)
+ return gv
+
+ example_args = (torch.arange(10, dtype=torch.float32) + 1,)
+ verify_model(Split6(), example_args, {}, Expected)
+
+ # Differential check against native PyTorch for non-divisible sizes and
dims.
+ class SplitModel(Module):
+ def __init__(self, split_size, dim):
+ super().__init__()
+ self.split_size = split_size
+ self.dim = dim
+
+ def forward(self, input):
+ return input.split(self.split_size, dim=self.dim)
+
+ def run_tvm(model, args):
+ exported_program = export(model, args=args)
+ mod = from_exported_program(exported_program)
+ ex = relax.build(mod, target="llvm")
+ vm = relax.VirtualMachine(ex, tvm.cpu())
+ out = vm["main"](*[tvm.runtime.tensor(a.numpy()) for a in args])
+ if hasattr(out, "numpy"):
+ return [out.numpy()]
+ return [o.numpy() for o in out]
+
+ for shape, split_size, dim in [
+ ((10,), 6, 0),
+ ((10,), 7, 0),
+ ((10,), 8, 0),
+ ((10,), 9, 0),
+ ((12,), 7, 0),
+ ((12, 8), 5, 1),
+ ((3, 10), 6, -1),
+ ]:
+ x = torch.arange(1, int(np.prod(shape)) + 1,
dtype=torch.float32).reshape(shape)
+ refs = [r.numpy() for r in x.split(split_size, dim)]
+ outs = run_tvm(SplitModel(split_size, dim), (x,))
+ assert [r.shape for r in refs] == [o.shape for o in outs], (
+ f"split shape={shape} s={split_size} dim={dim}: "
+ f"torch {[r.shape for r in refs]} vs tvm {[o.shape for o in outs]}"
+ )
+ for r, o in zip(refs, outs):
+ tvm.testing.assert_allclose(o, r, rtol=1e-7, atol=1e-7)
+
+
def test_squeeze():
class Squeeze1(Module):
def forward(self, input):
diff --git a/tests/python/relax/test_frontend_from_fx.py
b/tests/python/relax/test_frontend_from_fx.py
index a489977958..1d086cf043 100644
--- a/tests/python/relax/test_frontend_from_fx.py
+++ b/tests/python/relax/test_frontend_from_fx.py
@@ -4070,7 +4070,7 @@ def test_split():
R.Tensor((1, 1, 10, 10), dtype="float32"),
R.Tensor((1, 1, 10, 10), dtype="float32"),
R.Tensor((1, 1, 10, 10), dtype="float32"),
- ) = R.split(input_1, indices_or_sections=3, axis=1)
+ ) = R.split(input_1, indices_or_sections=[1, 2], axis=1)
gv: R.Tuple(
R.Tensor((1, 1, 10, 10), dtype="float32"),
R.Tensor((1, 1, 10, 10), dtype="float32"),
