_aten_to_core: fix IEEE-754 edge cases in replace_atan2

- Signed zeros: use 1/v trick (1/-0.0 = -inf) combined with strict >
  to make y_neg and x_neg correct for -0.0 inputs without misclassifying
  ±inf values which use the strict > path directly
- x = -0.0 branch: split x=0 into +0 (returns ±π/2) and -0 (returns ±π)
  per IEEE-754 atan2 specification
- Both infinite: add explicit branch for atan2(±inf, ±inf) → ±π/4 or
  ±3π/4; previously produced NaN via atan(inf/inf) = atan(NaN)
- Add test_signed_zeros and test_infinities to lock in correct behaviour
- Update IR tests to reflect the expanded op sequence

Author: gokulkrishna98

coreai_torch/_aten_to_core.py

@@ -1526,55 +1526,105 @@ def replace_atan2(values_map: dict[str, Value], node: fx.Node, loc: Location) ->
     """Lower atan2(y, x) using atan(y/x) with quadrant correction.
 
     CoreAI has no native atan2, so it is decomposed as:
-      - x != 0: atan(y/x) adjusted by ±π to place the result in the correct quadrant.
-      - x == 0: ±π/2 or 0 based on sign of y.
+      - x != 0, finite: atan(y/x) adjusted by ±π for the correct quadrant.
+      - x == +0: ±π/2 for non-zero y, 0 for y = 0.
+      - x == -0: ±π for all y (including ±0 → ±π per IEEE-754).
+      - both infinite: ±π/4 or ±3π/4 per IEEE-754.
+
+    Signed-zero handling: IEEE-754 treats -0.0 as distinct from +0.0 for atan2
+    (e.g. atan2(-0, -1) = -π, not +π). The 1/v trick — 1/-0.0 = -inf — is used
+    to detect the sign bit of zero inputs so that y_neg and x_neg are correct
+    for -0.0 inputs without misclassifying ±inf (which use the strict > path).
 
     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.
+    intermediate result is discarded by the final where-select.
     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
 
     zero = coreai.constant(0.0, dtype=ele_type)
+    one = coreai.constant(1.0, dtype=ele_type)
     pi = coreai.constant(np.pi, dtype=ele_type)
+    neg_pi = coreai.constant(-np.pi, dtype=ele_type)
     half_pi = coreai.constant(np.pi / 2.0, dtype=ele_type)
     neg_half_pi = coreai.constant(-np.pi / 2.0, dtype=ele_type)
-
-    # Avoid division by zero when x = 0 by substituting x = 1 for the ratio.
+    quarter_pi = coreai.constant(np.pi / 4.0, dtype=ele_type)
+    neg_quarter_pi = coreai.constant(-np.pi / 4.0, dtype=ele_type)
+    three_quarter_pi = coreai.constant(3.0 * np.pi / 4.0, dtype=ele_type)
+    neg_three_quarter_pi = coreai.constant(-3.0 * np.pi / 4.0, dtype=ele_type)
+
+    # ── signed-zero-aware sign predicates ─────────────────────────────────────
+    # 1 / -0.0 = -inf (IEEE-754), so (0 > 1/v) is True iff v = -0.0. Combine with
+    # the strict > predicate (handles ±inf and non-zero finites) via OR.
+    y_is_zero = coreai.broadcasting_equal(y, zero)
     x_is_zero = coreai.broadcasting_equal(x, zero)
-    x_safe = coreai.broadcasting_where(
-        x_is_zero, coreai.constant(1.0, dtype=ele_type), x
+    y_neg = coreai.broadcasting_or(
+        coreai.broadcasting_greater(zero, y),
+        coreai.broadcasting_and(
+            y_is_zero,
+            coreai.broadcasting_greater(zero, coreai.broadcasting_divide(one, y)),
+        ),
+    )
+    x_neg = coreai.broadcasting_or(
+        coreai.broadcasting_greater(zero, x),
+        coreai.broadcasting_and(
+            x_is_zero,
+            coreai.broadcasting_greater(zero, coreai.broadcasting_divide(one, x)),
+        ),
+    )
+    x_is_neg_zero = coreai.broadcasting_and(
+        x_is_zero,
+        coreai.broadcasting_greater(zero, coreai.broadcasting_divide(one, x)),
+    )
+
+    # ── both-infinite branch ──────────────────────────────────────────────────
+    # atan(inf/inf) = atan(NaN) = NaN; handle before the divide.
+    pos_inf = coreai.constant(float("inf"), dtype=ele_type)
+    neg_inf = coreai.constant(float("-inf"), dtype=ele_type)
+    x_is_inf = coreai.broadcasting_or(
+        coreai.broadcasting_equal(x, pos_inf), coreai.broadcasting_equal(x, neg_inf)
+    )
+    y_is_inf = coreai.broadcasting_or(
+        coreai.broadcasting_equal(y, pos_inf), coreai.broadcasting_equal(y, neg_inf)
+    )
+    both_inf = coreai.broadcasting_and(x_is_inf, y_is_inf)
+    inf_result = coreai.broadcasting_where(
+        y_neg,
+        coreai.broadcasting_where(x_neg, neg_three_quarter_pi, neg_quarter_pi),
+        coreai.broadcasting_where(x_neg, three_quarter_pi, quarter_pi),
     )
-    base = coreai.atan(coreai.broadcasting_divide(y, x_safe))
 
-    # Quadrant correction: x < 0 shifts the result by ±π.
-    x_neg = coreai.broadcasting_greater(zero, x)
-    y_neg = coreai.broadcasting_greater(zero, y)
-    y_pos = coreai.broadcasting_greater(y, zero)
+    # ── x = 0 branch ──────────────────────────────────────────────────────────
+    # x = +0: ±π/2 for strictly ±y, 0 when y = 0.
+    # x = -0: ±π for all y (y_neg covers y = -0.0 via the 1/y trick above).
+    y_pos_strict = coreai.broadcasting_greater(y, zero)
+    y_neg_strict = coreai.broadcasting_greater(zero, y)
+    pos_x_zero_result = coreai.broadcasting_where(
+        y_pos_strict,
+        half_pi,
+        coreai.broadcasting_where(y_neg_strict, neg_half_pi, zero),
+    )
+    neg_x_zero_result = coreai.broadcasting_where(y_neg, neg_pi, pi)
+    zero_result = coreai.broadcasting_where(
+        x_is_neg_zero, neg_x_zero_result, pos_x_zero_result
+    )
+
+    # ── finite nonzero x branch ────────────────────────────────────────────────
+    # Avoid division by zero: substitute x = 1 when x = 0; result discarded by
+    # the outer where-select.
+    x_safe = coreai.broadcasting_where(x_is_zero, one, x)
+    base = coreai.atan(coreai.broadcasting_divide(y, x_safe))
     correction = coreai.broadcasting_where(
         y_neg,
         coreai.broadcasting_sub(base, pi),
         coreai.broadcasting_add(base, pi),
     )
     nonzero_result = coreai.broadcasting_where(x_neg, correction, base)
 
-    # x = 0: result is π/2, −π/2, or 0 based on sign of y.
-    zero_result = coreai.broadcasting_where(
-        y_pos,
-        half_pi,
-        coreai.broadcasting_where(y_neg, neg_half_pi, zero),
-    )
-    return coreai.broadcasting_where(x_is_zero, zero_result, nonzero_result)
+    # ── combine ────────────────────────────────────────────────────────────────
+    result = coreai.broadcasting_where(x_is_zero, zero_result, nonzero_result)
+    return coreai.broadcasting_where(both_inf, inf_result, result)
 
 
 def replace_gather(values_map: dict[str, Value], node: fx.Node, loc: Location) -> Value:

tests/ops/test_ops.py

@@ -652,6 +652,22 @@ async def test_broadcast_shapes(self) -> None:
         x = torch.randn(4)
         await validate_numerical_output(model=model, y=y, x=x)
 
+    async def test_signed_zeros(self) -> None:
+        """IEEE-754 signed-zero cases: atan2(-0, x) and atan2(y, -0)."""
+        model = self.Atan2Model().eval()
+        # y = -0.0 with various x signs
+        y = torch.tensor([-0.0, -0.0, -0.0, 0.0])
+        x = torch.tensor([-1.0, 1.0, -0.0, -0.0])
+        await validate_numerical_output(model=model, y=y, x=x)
+
+    async def test_infinities(self) -> None:
+        """IEEE-754 both-infinite cases: atan2(±inf, ±inf) → ±π/4 or ±3π/4."""
+        model = self.Atan2Model().eval()
+        inf = float("inf")
+        y = torch.tensor([inf, inf, -inf, -inf])
+        x = torch.tensor([inf, -inf, inf, -inf])
+        await validate_numerical_output(model=model, y=y, x=x)
+
 
 @pytest.mark.parametrize(
     "x",