Skip to content

Commit 922176e

Browse files
address PR review feedback
- Document IEEE-754 limitations (signed zeros, infinities) in replace_atan2 docstring - Add (0, 0) case to test_x_zero to lock in atan2(0, 0) = 0 by convention Co-Authored-By: Claude Sonnet 4.6 (1M context) <noreply@anthropic.com>
1 parent 0dd1ccf commit 922176e

2 files changed

Lines changed: 12 additions & 3 deletions

File tree

coreai_torch/_aten_to_core.py

Lines changed: 9 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1532,6 +1532,15 @@ def replace_atan2(values_map: dict[str, Value], node: fx.Node, loc: Location) ->
15321532
When x=0, x is replaced with 1 before the divide solely to avoid NaN/inf; that
15331533
intermediate result is discarded by the final where-select in favour of the x=0 branch.
15341534
atan2(0, 0) = 0 by convention.
1535+
1536+
IEEE-754 limitations:
1537+
- Signed zeros: ``-0.0`` is treated the same as ``+0.0`` because the
1538+
``0 > y`` predicate is false for ``y = -0.0``. Results are numerically
1539+
equal to PyTorch for finite inputs but the sign bit may differ
1540+
(e.g. ``atan2(-0.0, -1.0)`` returns ``+π`` here, ``-π`` in PyTorch).
1541+
- Infinities: ``atan2(±inf, ±inf)`` returns NaN because ``inf/inf``
1542+
produces NaN before ``atan`` is applied. PyTorch returns ``±π/4``
1543+
or ``±3π/4`` per IEEE-754. Do not pass infinite inputs to this op.
15351544
"""
15361545
y, x = _get_operands(values_map, node, [0, 1])
15371546
ele_type = x.type.element_type

tests/ops/test_ops.py

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -626,10 +626,10 @@ async def test_basic(
626626
)
627627

628628
async def test_x_zero(self) -> None:
629-
"""x = 0 should yield ±π/2 depending on sign of y."""
629+
"""x = 0 should yield ±π/2 depending on sign of y; (0, 0) → 0 by convention."""
630630
model = self.Atan2Model().eval()
631-
y = torch.tensor([1.0, -1.0, 2.0, -2.0])
632-
x = torch.zeros(4)
631+
y = torch.tensor([1.0, -1.0, 2.0, -2.0, 0.0])
632+
x = torch.zeros(5)
633633
await validate_numerical_output(model=model, y=y, x=x)
634634

635635
async def test_y_zero(self) -> None:

0 commit comments

Comments
 (0)