Skip to content

Tell apart torch overloads that TorchScript reports under one node kind - #2809

Open
LeSingh1 wants to merge 1 commit into
apple:mainfrom
LeSingh1:torchscript-overload-arity
Open

Tell apart torch overloads that TorchScript reports under one node kind#2809
LeSingh1 wants to merge 1 commit into
apple:mainfrom
LeSingh1:torchscript-overload-arity

Conversation

@LeSingh1

@LeSingh1 LeSingh1 commented Aug 9, 2026

Copy link
Copy Markdown
Contributor

TorchScript drops the overload suffix, so several aten ops reach the converter as a single node kind with two different argument layouts. Three converters assume the wrong one. In each case the overloads differ in arity, which is what the fix keys on.

call traced as today
torch.sum(x, dtype=torch.float32) aten::sum, 2 inputs IndexError: pop index out of range
torch.mean(x, dtype=torch.float32) aten::mean, 2 inputs same
torch.randint(10, (2, 3)) aten::randint, 6 inputs TypeError: 'int' object is not iterable
torch.round(x, decimals=2) aten::round, 2 inputs ValueError: node 5 (round) got 2 input(s), expected [1]

sum / meanaten::sum(self, *, dtype) has 2 inputs; aten::sum.dim_IntList(self, dim, keepdim, dtype) has 4. The converter always read position 1 as dim, so the dtype enum (6 for fp32) was passed as an axis. The torch.export branch already discriminates by node kind; only the TorchScript branch was missing it.

randintaten::randint(high, size, ...) has 6 inputs; aten::randint.low(low, high, size, ...) has 7. Under TorchScript the converter unconditionally took the .low layout, reading the size list as high and the dtype enum as the shape.

roundaten::round.decimals(self, decimals) was rejected outright by expected=1. Accepted now, and decimals is honored by scaling by 10**decimals; dropping it would silently change the result. decimals is also read from keyword inputs, for the export path.

Verified arities by tracing, e.g. torch.sum(x) and torch.sum(x, dtype=...) are both 2 inputs while torch.sum(x, dim=0), torch.sum(x, 0, True) and torch.sum(x, dim=0, dtype=...) are all 4, so arity separates them cleanly. logsumexp (3), all/any (1 or 3) are unaffected.

Testing

TestSum::test_sum_mean_dtype, TestRandint::test_randint_no_low, TestElementWiseUnary::test_round_decimals (decimals 0, 1, 2, -1). All 13 TorchScript cases fail on main and pass here. Existing TestSum, TestRandint and TestElementWiseUnary tests still pass.

This machine has a broken scikit-learn install that makes every TorchFrontend.TORCHEXPORT case error out, including pre-existing ones such as TestElementWiseUnary::test_acosh, so I could only exercise the TorchScript frontend locally.

TorchScript drops the overload suffix, so several aten ops reach the
converter as one kind with two different argument layouts. Three of them
assume the wrong layout. In each case the two overloads have different
arity, which is what the fix keys on.

sum / mean: aten::sum(self, dtype) has 2 inputs, aten::sum.dim_IntList(
self, dim, keepdim, dtype) has 4. The converter always read position 1 as
dim, so torch.sum(x, dtype=torch.float32) passed the dtype enum as an axis
and failed with "IndexError: pop index out of range".

randint: aten::randint(high, size, ...) has 6 inputs, aten::randint.low(
low, high, size, ...) has 7. Under TorchScript the converter always took
the .low layout, so torch.randint(10, (2, 3)) read the size list as high
and the dtype enum as the shape, failing with "TypeError: 'int' object is
not iterable".

round: aten::round.decimals(self, decimals) has 2 inputs and was rejected
outright by expected=1. Accept it and scale by 10**decimals, since
dropping the argument would silently change the result.
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant