Skip to content

[Fix][Relax][Frontend][Torch] Accumulate integer cumsum and cumprod in int64 - #20374

Open
hiyufan wants to merge 1 commit into
apache:mainfrom
hiyufan:fix/relax-torch-cumsum-int64
Open

hiyufan wants to merge 1 commit into
apache:mainfrom
hiyufan:fix/relax-torch-cumsum-int64

Conversation

@hiyufan

@hiyufan hiyufan commented Sep 17, 2026

Copy link
Copy Markdown
Contributor

Problem

With no dtype argument, torch.cumsum / torch.cumprod accumulate every integral and bool input in int64. The converters passed dtype=None through to relax.op.cumsum / cumprod, so the running sum kept the input dtype — and wrapped:

input torch frontend on main
uint8 [200, 100, 50].cumsum(1) int64 [200, 300, 350] uint8 [200, 44, 94]
int8 [100, 100, 50].cumsum(1) int64 [100, 200, 250] int8 [100, -56, -6]
int32 [2³⁰, 2³⁰, 5].cumsum(1) int64 […, 2147483653] int32 […, -2147483643]
uint8 [200, 100, 50].cumprod(1) int64 [200, 20000, 1000000] uint8 [200, 32, 64]
bool [T, T, F].cumsum(0) int64 InternalError (relax cumsum on bool)

Not a dtype-label difference: the values are wrong as soon as the sum leaves the narrow range, which for uint8/int8 is the second element.

Fix

_cumulative_dtype returns the explicit dtype= if given, otherwise int64 for an integral or bool input and None (keep) for floats — torch's rule. Both converters use it.

Verification

cumsum and cumprod along both axes plus cumsum(dtype=float32), over bool, uint8, int8, int16, int32, int64, float16, float32 and float64 inputs chosen to overflow the narrow types, each built with relax.build(llvm) and compared with torch on dtype and values:

matched wrong raised
main 25 16 4
this PR 45 0 0

The explicit-dtype variant and every float input were already correct and are unchanged. test_frontend_from_exported_program.py and test_frontend_from_fx.py: failure sets identical before and after apart from the new tests. ruff check / ruff format --check (v0.12.3, the version CI pins) clean.

Tests

  • test_cumsum_integer_input_accumulates_in_int64 — IR-level: a uint8 input emits R.cumsum(x, axis=1, dtype="int64").
  • test_cumsum_cumprod_integer_valuescumsum, cumprod and cumsum(dtype=float32) over bool, uint8, int8, int32 and int64, asserting dtype and values.

Against the previous head the bool, uint8, int8 and int32 rows fail (assert 'uint8' == 'int64' and the IR mismatch); int64 passes there and pins that it is untouched.

Independent of #20372 / #20373 (branched from main).


This change was prepared with AI assistance (Claude). I have reviewed and verified it, and can speak to it in review.

…n int64

With no `dtype` argument torch accumulates every integral and bool input of
`cumsum` / `cumprod` in int64. The converters passed `dtype=None` through, so the
running sum kept the input dtype and wrapped:

  uint8 [200, 100, 50].cumsum(1)     torch int64 [200, 300, 350]        before uint8 [200, 44, 94]
  int8  [100, 100, 50].cumsum(1)     torch int64 [100, 200, 250]        before int8  [100, -56, -6]
  int32 [2^30, 2^30, 5].cumsum(1)    torch int64 [.., 2147483653]       before int32 [.., -2147483643]
  bool  [T, T, F].cumsum(0)          torch int64 [1, 1, 0, ...]         before InternalError

An explicit `dtype=` was already honoured and is unchanged; float inputs keep their
dtype, as in torch. The two converters share the rule through `_cumulative_dtype`.

Swept cumsum / cumprod along both axes plus cumsum(dtype=float32) over bool, uint8,
int8, int16, int32, int64, float16, float32 and float64 inputs chosen to overflow the
narrow types, built with relax.build(llvm) and compared with torch on dtype and
values: 25 matched / 16 wrong / 4 raised before, 45 / 0 / 0 after.

Tests: an IR-level check that a uint8 cumsum emits `R.cumsum(..., dtype="int64")`,
and numeric checks of cumsum, cumprod and cumsum(dtype=float32) over bool, uint8,
int8, int32 and int64 inputs. The bool, uint8, int8 and int32 cases fail against the
previous head; int64 passes there and pins that it is untouched.
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