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

Author: gokulkrishna98

coreai_torch/_aten_to_core.py

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

tests/ops/test_ops.py

@@ -626,10 +626,10 @@ async def test_basic(
         )
 
     async def test_x_zero(self) -> None:
-        """x = 0 should yield ±π/2 depending on sign of y."""
+        """x = 0 should yield ±π/2 depending on sign of y; (0, 0) → 0 by convention."""
         model = self.Atan2Model().eval()
-        y = torch.tensor([1.0, -1.0, 2.0, -2.0])
-        x = torch.zeros(4)
+        y = torch.tensor([1.0, -1.0, 2.0, -2.0, 0.0])
+        x = torch.zeros(5)
         await validate_numerical_output(model=model, y=y, x=x)
 
     async def test_y_zero(self) -> None: