siyiweigeHEW opened a new issue, #20232:
URL: https://github.com/apache/tvm/issues/20232
### Expected behavior
A valid PyTorch model calling `x.split(split_size, dim)` must produce the
same chunk shapes as PyTorch: chunks of `split_size` elements along `dim`, with
the last chunk smaller if `D % split_size != 0`. For example, for `x` with
shape `(10,)`:
- `x.split(6)` → two chunks of shapes `(6,)` and `(4,)` (values `[1..6]` and
`[7..10]`)
- `x.split(7)` → `(7,)` and `(3,)`
- `x.split(3)` → `(3,), (3,), (3,), (1,)`
The model passes `torch.export.export` and runs correctly in PyTorch.
### Actual behavior
`tvm.relax.frontend.torch.from_exported_program` produces wrong chunk shapes
whenever the input dimension is not divisible by `split_size` and `split_size >
D/2` (more generally, whenever `ceil(D / ceil(D / split_size)) != split_size`).
The converter `_split` in
`python/tvm/relax/frontend/torch/base_fx_graph_translator.py:2316-2327` (mapped
from `aten.split.Tensor` and `aten.split_with_sizes.default` in
`exported_program_translator.py:1959-1960`) converts the per-chunk size to a
**section count**:
```python
else: # int split_size
n_section = (self.shape_of(x)[dim].value + split_size - 1) // split_size
return self.block_builder.emit(relax.op.split(x, n_section, dim))
```
and passes it to `relax.op.split`'s integer argument. But `relax.op.split`'s
integer argument is **the number of (equal) sections** (`split_len = ceildiv(D,
n_section)`, `src/relax/op/tensor/manipulate.cc:1142-1169`), not the per-chunk
size. For `D=10, split_size=6`, the frontend emits `R.split(x,
indices_or_sections=2, axis=0)`, which yields two chunks of `ceildiv(10, 2) =
5` elements — `(5,)` and `(5,)` — instead of PyTorch's `(6,)` and `(4,)`.
For `D=10`:
| `split_size` | PyTorch | TVM (from_exported_program) |
|---|---|---|
| 6 | `(6,), (4,)` | `(5,), (5,)` |
| 7 | `(7,), (3,)` | `(5,), (5,)` |
| 8 | `(8,), (2,)` | `(5,), (5,)` |
| 9 | `(9,), (1,)` | `(5,), (5,)` |
| 3 | `(3,), (3,), (3,), (1,)` | `(3,), (3,), (3,), (1,)` (coincidentally
matches) |
The bug is not limited to `dim=0`: e.g. `(12, 8)`, `x.split(5, dim=1)` →
PyTorch `(12,5), (12,3)` vs TVM `(12,4), (12,4)`.
The `list/tuple` form (`x.split([s0, s1, ...])`,
`aten.split_with_sizes.default`) is handled correctly: cumulative indices `[s0,
s0+s1, ...]` are passed to `relax.op.split`.
### Environment
- OS: Linux
- TVM: v0.24.dev0 (main branch; fuzzed commit `262c6d2e0`; current main
commit `390af87345` has byte-identical `_split` code — the bug persists)
- Python: 3.11 (fuzzed) / 3.10 (verification env)
- torch: 2.6.0
### Steps to reproduce
```python
"""Repro: TVM relax torch frontend mishandles x.split(int) with split_size >
D/2."""
import numpy as np
import torch
import torch.nn as nn
import tvm
from tvm import relax
from tvm.relax.frontend.torch import from_exported_program
class SplitModule(nn.Module):
def __init__(self, split_size):
super().__init__()
self.split_size = split_size
def forward(self, x):
return x.split(self.split_size, dim=0)
def run_tvm(x, split_size):
m = SplitModule(split_size).float().eval()
exp = torch.export.export(m, (x,))
mod = from_exported_program(exp) # <- current frontend,
_split code above
ex = relax.build(mod, target="llvm")
vm = relax.VirtualMachine(ex, tvm.cpu())
out = vm["main"](tvm.nd.array(x.numpy()))
return [o.numpy() for o in out]
x = torch.arange(10, dtype=torch.float32) + 1 # [1..10]
refs = x.split(6) # torch: (6,), (4,)
print("torch:", [r.shape for r in refs], [r.numpy().tolist() for r in refs])
outs = run_tvm(x, 6)
print("tvm: ", [o.shape for o in outs], [o.tolist() for o in outs])
```
Actual output:
```
torch: [torch.Size([6]), torch.Size([4])] [[1.0, 2.0, 3.0, 4.0, 5.0, 6.0],
[7.0, 8.0, 9.0, 10.0]]
tvm: [(5,), (5,)] [[1.0, 2.0, 3.0, 4.0, 5.0], [6.0, 7.0, 8.0, 9.0, 10.0]]
```
The IR produced by `from_exported_program` for `x.split(6)` on `x:
R.Tensor((10,))`:
```
lv: R.Tuple(R.Tensor((5,), dtype="float32"), R.Tensor((5,),
dtype="float32")) =
R.split(x, indices_or_sections=2, axis=0) # PyTorch wants chunks of
6: (6,), (4,)
```
### Root cause
In `_split` (`base_fx_graph_translator.py:2316-2327`), the int `split_size`
branch computes `n_section = ceil(D / split_size)` and passes it as
`relax.op.split`'s integer `indices_or_sections`, which relax interprets as
"split into `n_section` equal parts of size `ceil(D / n_section)`". PyTorch's
`x.split(s)` means "chunks of size `s`". These coincide only when `ceil(D /
ceil(D / s)) == s` (e.g. divisible sizes, or `D=10, s=3`), so valid models
silently produce differently-shaped tensors.
Suggested fix: for int `split_size`, pass the cumulative cut positions
instead of the section count, e.g. `indices = [s, 2*s, ..., (ceil(D/s) - 1) *
s]` (the same cumulative-index list the `list/tuple` branch already builds).
### Triage
* needs-triage
* bug
* relax
* frontend/torch
--
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]
