Skip to content

Commit efcc690

Browse files
_aten_to_core: gate si32 arange path on all operands being integer-typed
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.
1 parent e8dac1d commit efcc690

2 files changed

Lines changed: 21 additions & 7 deletions

File tree

coreai_torch/_aten_to_core.py

Lines changed: 7 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -644,15 +644,16 @@ def replace_arange_start_step(
644644
else coreai.constant(1, dtype=start.type.element_type)
645645
)
646646

647-
# When operands are integer-typed, keep them as si32 so coreai.range_ can
648-
# infer a static output shape, then cast the result to the requested dtype.
649-
# For float operands the count isn't statically determinable anyway, so
650-
# cast everything to target_type and let range_ return a dynamic shape.
647+
# When ALL operands are integer-typed, keep them as si32 so coreai.range_
648+
# can infer a static output shape, then cast the result to the requested
649+
# dtype. If any operand is float (e.g. arange(0, 5, 0.5)), fall back to
650+
# target_type — truncating a float step to si32 would corrupt the values.
651651
target_type = get_output_element_type_from_node(node)
652652
si32 = IntegerType.get_signed(32)
653-
range_type = (
654-
si32 if isinstance(start.type.element_type, IntegerType) else target_type
653+
all_integer = all(
654+
isinstance(v.type.element_type, IntegerType) for v in (start, end, step)
655655
)
656+
range_type = si32 if all_integer else target_type
656657

657658
def to_scalar(v: Value) -> Value:
658659
if v.type.rank > 0:

tests/ops/test_ops.py

Lines changed: 14 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -347,7 +347,20 @@ def forward(self, x: Tensor) -> Tensor:
347347
x = torch.zeros(2, 8)
348348
await validate_numerical_output(model=model, x=x)
349349

350-
async def test_symint_end_with_float_start_step(self) -> None:
350+
async def test_mixed_int_float_operands(self) -> None:
351+
"""Regression: arange with mixed int/float operands must not truncate.
352+
353+
torch.arange(0, 5, 0.5) has int start/end but float step. The lowering
354+
must not cast step to si32 (which would truncate 0.5 → 0 and produce a
355+
degenerate range); it must fall back to the float path instead.
356+
"""
357+
358+
class ArangeMixedScalars(nn.Module):
359+
def forward(self) -> Tensor:
360+
return torch.arange(0, 5, 0.5, dtype=torch.float32)
361+
362+
await validate_numerical_output(model=ArangeMixedScalars().eval())
363+
351364
"""Regression for ``replace_arange_start_step``: when ``end`` is
352365
SymInt-derived (carrying f32 element type from a sym_size cast)
353366
and ``start`` / ``step`` come in as scalar si32, ``coreai.range_``

0 commit comments

Comments
 (0)