diff --git a/coreai_torch/_aten_to_core.py b/coreai_torch/_aten_to_core.py index 20a2f83..e589436 100644 --- a/coreai_torch/_aten_to_core.py +++ b/coreai_torch/_aten_to_core.py @@ -644,15 +644,16 @@ def replace_arange_start_step( else coreai.constant(1, dtype=start.type.element_type) ) - # When operands are integer-typed, keep them as si32 so coreai.range_ can - # infer a static output shape, then cast the result to the requested dtype. - # For float operands the count isn't statically determinable anyway, so - # cast everything to target_type and let range_ return a dynamic shape. + # When ALL operands are integer-typed, keep them as si32 so coreai.range_ + # can infer a static output shape, then cast the result to the requested + # dtype. If any operand is float (e.g. arange(0, 5, 0.5)), fall back to + # target_type — truncating a float step to si32 would corrupt the values. target_type = get_output_element_type_from_node(node) si32 = IntegerType.get_signed(32) - range_type = ( - si32 if isinstance(start.type.element_type, IntegerType) else target_type + all_integer = all( + isinstance(v.type.element_type, IntegerType) for v in (start, end, step) ) + range_type = si32 if all_integer else target_type def to_scalar(v: Value) -> Value: if v.type.rank > 0: diff --git a/tests/ops/test_ops.py b/tests/ops/test_ops.py index 6e85b7b..e244d28 100644 --- a/tests/ops/test_ops.py +++ b/tests/ops/test_ops.py @@ -347,7 +347,20 @@ def forward(self, x: Tensor) -> Tensor: x = torch.zeros(2, 8) await validate_numerical_output(model=model, x=x) - async def test_symint_end_with_float_start_step(self) -> None: + async def test_mixed_int_float_operands(self) -> None: + """Regression: arange with mixed int/float operands must not truncate. + + torch.arange(0, 5, 0.5) has int start/end but float step. The lowering + must not cast step to si32 (which would truncate 0.5 → 0 and produce a + degenerate range); it must fall back to the float path instead. + """ + + class ArangeMixedScalars(nn.Module): + def forward(self) -> Tensor: + return torch.arange(0, 5, 0.5, dtype=torch.float32) + + await validate_numerical_output(model=ArangeMixedScalars().eval()) + """Regression for ``replace_arange_start_step``: when ``end`` is SymInt-derived (carrying f32 element type from a sym_size cast) and ``start`` / ``step`` come in as scalar si32, ``coreai.range_``