[Bug][Relax][Frontend][PyTorch] `x.split(split_size)` with a non-divisible `split_size > D/2` produces wrong chunk shapes: frontend `_split` converts the per-chunk size to a section count, but `relax.op.split`'s int argument means "equal sections"

[siyiweigeHEW]

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:

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

"""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