Skip to content

fix(arange): guard si32 path on all operands being integer-typed#27

Merged
gokulkrishna98 merged 1 commit into
apple:mainfrom
gokulkrishna98:dev/gokul/fix-arange-mixed-operands
Jun 29, 2026
Merged

fix(arange): guard si32 path on all operands being integer-typed#27
gokulkrishna98 merged 1 commit into
apple:mainfrom
gokulkrishna98:dev/gokul/fix-arange-mixed-operands

Conversation

@gokulkrishna98

Copy link
Copy Markdown
Contributor

Problem

torch.arange accepts mixed int/float scalar bounds — e.g. arange(0, 5, 0.5) where start/end are int but step is float. The previous fix gated the si32 optimisation on start.type.element_type alone, so step (arriving as f32) would be cast to si32, truncating 0.5 → 0 and producing a degenerate range.

Fix

Check all three operands before selecting the integer path:

all_integer = all(isinstance(v.type.element_type, IntegerType) for v in (start, end, step))
range_type = si32 if all_integer else target_type

If any operand is float, fall back to target_type (float) for the range_ call, preserving the original behaviour for float-step aranges.

Tests

  • Added TestArange::test_mixed_int_float_operands: arange(0, 5, 0.5, dtype=float32) must produce the correct 10-element output and not truncate the float step.

Checking only start.type.element_type was insufficient: mixed-type calls
like arange(0, 5, 0.5) have int start/end but float step, so the old
guard would select si32 and truncate step (0.5 → 0), producing a
degenerate range. Now all three operands must be integer-typed for the
static-shape si32 path; any float operand falls back to target_type.

Adds test_mixed_int_float_operands to cover this case.
@gokulkrishna98 gokulkrishna98 force-pushed the dev/gokul/fix-arange-mixed-operands branch from b7da02f to efcc690 Compare June 29, 2026 22:48
@gokulkrishna98 gokulkrishna98 merged commit 779df85 into apple:main Jun 29, 2026
2 checks passed
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.

2 participants