Skip to content

Commit 1b3cb3b

Browse files
gokulkrishna98gokulkrishna98
authored
atan2: propagate NaN instead of returning 0/pi at x == +/-0 (#35)
## Description The `atan2` converter's x=0 branch classifies inputs purely via comparisons (`y > 0`, `0 > y`, etc.), which are all `False` when NaN is involved. This meant `atan2(NaN, +0.0)` incorrectly returned `0` and `atan2(NaN, -0.0)` incorrectly returned `pi`, instead of `NaN` as required by IEEE-754 and matched by `torch.atan2`. This adds an explicit NaN check (`y != y or x != x`) that takes priority over every other branch in the decomposition, so NaN now always propagates through correctly. ## Testing - Added `test_nan_propagation` to `TestAtan2` in `tests/ops/test_ops.py`, covering NaN in `y`, `x`, and both, including at `x = +/-0`. - Updated the three `TestAtan2IR` FileCheck patterns in `tests/ops/test_ops_ir.py` to match the new IR shape (extra NaN constant + final NaN-select `where`). - Full `tests/ops/` suite passes locally.
1 parent 40312b5 commit 1b3cb3b

3 files changed

Lines changed: 39 additions & 4 deletions

File tree

‎coreai_torch/_aten_to_core.py‎

Lines changed: 14 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1555,6 +1555,9 @@ def replace_atan2(values_map: dict[str, Value], node: fx.Node, loc: Location) ->
15551555
- x == +0: ±π/2 for non-zero y, 0 for y = 0.
15561556
- x == -0: ±π for all y (including ±0 → ±π per IEEE-754).
15571557
- both infinite: ±π/4 or ±3π/4 per IEEE-754.
1558+
- either operand is NaN: NaN, checked last so it overrides every other
1559+
branch (comparisons against NaN are all False, which would otherwise
1560+
misclassify NaN as one of the zero/quadrant cases above).
15581561
15591562
Signed-zero handling: IEEE-754 treats -0.0 as distinct from +0.0 for atan2
15601563
(e.g. atan2(-0, -1) = -π, not +π). The 1/v trick — 1/-0.0 = -inf — is used
@@ -1603,6 +1606,15 @@ def replace_atan2(values_map: dict[str, Value], node: fx.Node, loc: Location) ->
16031606
coreai.broadcasting_greater(zero, coreai.broadcasting_divide(one, x)),
16041607
)
16051608

1609+
# ── NaN branch ─────────────────────────────────────────────────────────────
1610+
# NaN != NaN under IEEE-754, so this is a self-contained NaN check. Needed
1611+
# because the x=0 branch below classifies purely on comparisons, which are
1612+
# all False for NaN and would otherwise misclassify atan2(NaN, ±0).
1613+
any_nan = coreai.broadcasting_or(
1614+
coreai.broadcasting_not_equal(y, y), coreai.broadcasting_not_equal(x, x)
1615+
)
1616+
nan_result = coreai.constant(float("nan"), dtype=ele_type)
1617+
16061618
# ── both-infinite branch ──────────────────────────────────────────────────
16071619
# atan(inf/inf) = atan(NaN) = NaN; handle before the divide.
16081620
pos_inf = coreai.constant(float("inf"), dtype=ele_type)
@@ -1649,7 +1661,8 @@ def replace_atan2(values_map: dict[str, Value], node: fx.Node, loc: Location) ->
16491661

16501662
# ── combine ────────────────────────────────────────────────────────────────
16511663
result = coreai.broadcasting_where(x_is_zero, zero_result, nonzero_result)
1652-
return coreai.broadcasting_where(both_inf, inf_result, result)
1664+
result = coreai.broadcasting_where(both_inf, inf_result, result)
1665+
return coreai.broadcasting_where(any_nan, nan_result, result)
16531666

16541667

16551668
def replace_gather(values_map: dict[str, Value], node: fx.Node, loc: Location) -> Value:

‎tests/ops/test_ops.py‎

Lines changed: 10 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -681,6 +681,16 @@ async def test_infinities(self) -> None:
681681
x = torch.tensor([inf, -inf, inf, -inf])
682682
await validate_numerical_output(model=model, y=y, x=x)
683683

684+
async def test_nan_propagation(self) -> None:
685+
"""NaN in either operand must propagate to NaN, including at x = ±0
686+
where comparison-based branch selection would otherwise misclassify
687+
NaN as one of the zero/quadrant cases."""
688+
model = self.Atan2Model().eval()
689+
nan = float("nan")
690+
y = torch.tensor([nan, nan, 1.0, nan, 0.0])
691+
x = torch.tensor([0.0, -0.0, nan, nan, nan])
692+
await validate_numerical_output(model=model, y=y, x=x)
693+
684694

685695
@pytest.mark.parametrize(
686696
"x",

‎tests/ops/test_ops_ir.py‎

Lines changed: 15 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -1175,6 +1175,7 @@ def forward(self, y: Tensor, x: Tensor) -> Tensor:
11751175
// CHECK-NEXT: coreai.graph @main(%[[Y:.*]]: tensor<2x3xf32> {coreai.name = "y"}, %[[X:.*]]: tensor<2x3xf32> {coreai.name = "x"}) -> (tensor<2x3xf32> {coreai.name = "{{.*}}"}) attributes {__coreai_pure__} {
11761176
// CHECK: %[[NEG_INF:.*]] = coreai.constant dense<0xFF800000> : tensor<f32>
11771177
// CHECK: %[[POS_INF:.*]] = coreai.constant dense<0x7F800000> : tensor<f32>
1178+
// CHECK: %[[NAN:.*]] = coreai.constant dense<0x7FC00000> : tensor<f32>
11781179
// CHECK: %[[ZERO:.*]] = coreai.constant dense<0.000000e+00> : tensor<f32>
11791180
// CHECK: %[[ONE:.*]] = coreai.constant dense<1.000000e+00> : tensor<f32>
11801181
// CHECK: %[[PI:.*]] = coreai.constant dense<3.14159274> : tensor<f32>
@@ -1197,6 +1198,9 @@ def forward(self, y: Tensor, x: Tensor) -> Tensor:
11971198
// CHECK: %[[RECIP_X_NEG:.*]] = coreai.decomposable.broadcasting_greater %[[ZERO]], %[[RECIP_X]]
11981199
// CHECK: %[[X_ZERO_NEG:.*]] = coreai.decomposable.broadcasting_and %[[X_IS_ZERO]], %[[RECIP_X_NEG]]
11991200
// CHECK: %[[X_NEG:.*]] = coreai.decomposable.broadcasting_or %[[X_NEG_STRICT]], %[[X_ZERO_NEG]]
1201+
// CHECK: %[[Y_NOT_EQ_Y:.*]] = coreai.decomposable.broadcasting_not_equal %[[Y]], %[[Y]]
1202+
// CHECK: %[[X_NOT_EQ_X:.*]] = coreai.decomposable.broadcasting_not_equal %[[X]], %[[X]]
1203+
// CHECK: %[[ANY_NAN:.*]] = coreai.decomposable.broadcasting_or %[[Y_NOT_EQ_Y]], %[[X_NOT_EQ_X]]
12001204
// CHECK: %[[X_IS_POS_INF:.*]] = coreai.decomposable.broadcasting_equal %[[X]], %[[POS_INF]]
12011205
// CHECK: %[[X_IS_NEG_INF:.*]] = coreai.decomposable.broadcasting_equal %[[X]], %[[NEG_INF]]
12021206
// CHECK: %[[X_IS_INF:.*]] = coreai.decomposable.broadcasting_or %[[X_IS_POS_INF]], %[[X_IS_NEG_INF]]
@@ -1205,7 +1209,8 @@ def forward(self, y: Tensor, x: Tensor) -> Tensor:
12051209
// CHECK: %[[Y_IS_INF:.*]] = coreai.decomposable.broadcasting_or %[[Y_IS_POS_INF]], %[[Y_IS_NEG_INF]]
12061210
// CHECK: %[[BOTH_INF:.*]] = coreai.decomposable.broadcasting_and %[[X_IS_INF]], %[[Y_IS_INF]]
12071211
// CHECK: %[[BASE:.*]] = coreai.atan
1208-
// CHECK: %[[RESULT:.*]] = coreai.decomposable.broadcasting_where %[[BOTH_INF]],
1212+
// CHECK: %[[INF_RESULT:.*]] = coreai.decomposable.broadcasting_where %[[BOTH_INF]],
1213+
// CHECK: %[[RESULT:.*]] = coreai.decomposable.broadcasting_where %[[ANY_NAN]], %[[NAN]], %[[INF_RESULT]]
12091214
// CHECK-NEXT: coreai.output %[[RESULT]] : tensor<2x3xf32>
12101215
// CHECK-NEXT: }
12111216
// CHECK-NEXT: }
@@ -1232,6 +1237,7 @@ def forward(self, y: Tensor, x: Tensor) -> Tensor:
12321237
// CHECK-NEXT: coreai.graph @main(%[[Y:.*]]: tensor<?x?xf32> {coreai.name = "y"}, %[[X:.*]]: tensor<?x?xf32> {coreai.name = "x"}) -> (tensor<?x?xf32> {coreai.name = "{{.*}}"}) attributes {__coreai_pure__} {
12331238
// CHECK: %[[NEG_INF:.*]] = coreai.constant dense<0xFF800000> : tensor<f32>
12341239
// CHECK: %[[POS_INF:.*]] = coreai.constant dense<0x7F800000> : tensor<f32>
1240+
// CHECK: %[[NAN:.*]] = coreai.constant dense<0x7FC00000> : tensor<f32>
12351241
// CHECK: %[[ZERO:.*]] = coreai.constant dense<0.000000e+00> : tensor<f32>
12361242
// CHECK: %[[ONE:.*]] = coreai.constant dense<1.000000e+00> : tensor<f32>
12371243
// CHECK: %[[PI:.*]] = coreai.constant dense<3.14159274> : tensor<f32>
@@ -1254,6 +1260,9 @@ def forward(self, y: Tensor, x: Tensor) -> Tensor:
12541260
// CHECK: %[[RECIP_X_NEG:.*]] = coreai.decomposable.broadcasting_greater %[[ZERO]], %[[RECIP_X]]
12551261
// CHECK: %[[X_ZERO_NEG:.*]] = coreai.decomposable.broadcasting_and %[[X_IS_ZERO]], %[[RECIP_X_NEG]]
12561262
// CHECK: %[[X_NEG:.*]] = coreai.decomposable.broadcasting_or %[[X_NEG_STRICT]], %[[X_ZERO_NEG]]
1263+
// CHECK: %[[Y_NOT_EQ_Y:.*]] = coreai.decomposable.broadcasting_not_equal %[[Y]], %[[Y]]
1264+
// CHECK: %[[X_NOT_EQ_X:.*]] = coreai.decomposable.broadcasting_not_equal %[[X]], %[[X]]
1265+
// CHECK: %[[ANY_NAN:.*]] = coreai.decomposable.broadcasting_or %[[Y_NOT_EQ_Y]], %[[X_NOT_EQ_X]]
12571266
// CHECK: %[[X_IS_POS_INF:.*]] = coreai.decomposable.broadcasting_equal %[[X]], %[[POS_INF]]
12581267
// CHECK: %[[X_IS_NEG_INF:.*]] = coreai.decomposable.broadcasting_equal %[[X]], %[[NEG_INF]]
12591268
// CHECK: %[[X_IS_INF:.*]] = coreai.decomposable.broadcasting_or %[[X_IS_POS_INF]], %[[X_IS_NEG_INF]]
@@ -1262,7 +1271,8 @@ def forward(self, y: Tensor, x: Tensor) -> Tensor:
12621271
// CHECK: %[[Y_IS_INF:.*]] = coreai.decomposable.broadcasting_or %[[Y_IS_POS_INF]], %[[Y_IS_NEG_INF]]
12631272
// CHECK: %[[BOTH_INF:.*]] = coreai.decomposable.broadcasting_and %[[X_IS_INF]], %[[Y_IS_INF]]
12641273
// CHECK: %[[BASE:.*]] = coreai.atan
1265-
// CHECK: %[[RESULT:.*]] = coreai.decomposable.broadcasting_where %[[BOTH_INF]],
1274+
// CHECK: %[[INF_RESULT:.*]] = coreai.decomposable.broadcasting_where %[[BOTH_INF]],
1275+
// CHECK: %[[RESULT:.*]] = coreai.decomposable.broadcasting_where %[[ANY_NAN]], %[[NAN]], %[[INF_RESULT]]
12661276
// CHECK-NEXT: coreai.output %[[RESULT]] : tensor<?x?xf32>
12671277
// CHECK-NEXT: }
12681278
// CHECK-NEXT: }
@@ -1284,9 +1294,11 @@ def forward(self, y: Tensor, x: Tensor) -> Tensor:
12841294
// CHECK: %[[ONE:.*]] = coreai.constant dense<1.000000e+00> : tensor<f32>
12851295
// CHECK: %[[Y_NEG:.*]] = coreai.decomposable.broadcasting_or
12861296
// CHECK: %[[X_NEG:.*]] = coreai.decomposable.broadcasting_or
1297+
// CHECK: %[[ANY_NAN:.*]] = coreai.decomposable.broadcasting_or
12871298
// CHECK: %[[BOTH_INF:.*]] = coreai.decomposable.broadcasting_and
12881299
// CHECK: %[[BASE:.*]] = coreai.atan
1289-
// CHECK: %[[RESULT:.*]] = coreai.decomposable.broadcasting_where %[[BOTH_INF]],
1300+
// CHECK: %[[INF_RESULT:.*]] = coreai.decomposable.broadcasting_where %[[BOTH_INF]],
1301+
// CHECK: %[[RESULT:.*]] = coreai.decomposable.broadcasting_where %[[ANY_NAN]],
12901302
// CHECK-NEXT: coreai.output %[[RESULT]] : tensor<4xf32>
12911303
// CHECK-NEXT: }
12921304
// CHECK-NEXT: }

0 commit comments

Comments
 (0)