siyiweigeHEW opened a new pull request, #20240:
URL: https://github.com/apache/tvm/pull/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.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
