Skip to content

[converter] implement aten::atan2 conversion #23

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
gokulkrishna98 merged 8 commits into from
Jul 8, 2026
+313 −0
Merged

[converter] implement aten::atan2 conversion #23

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
8 commits
Select commit Hold shift + click to select a range
a79f21e
_aten_to_core: implement aten::atan2
gokulkrishna98 Jun 22, 2026
4b9c442
_aten_to_core: expand replace_atan2 docstring
gokulkrishna98 Jun 22, 2026
ba24874
tests/ops: restore missing x parametrize for test_batchnorm
gokulkrishna98 Jun 22, 2026
0dd1ccf
docs: strip notebook outputs (nbstripout)
gokulkrishna98 Jun 23, 2026
ab8cdb3
address PR review feedback
gokulkrishna98 Jun 23, 2026
5e2a2d7
_aten_to_core: fix IEEE-754 edge cases in replace_atan2
gokulkrishna98 Jun 23, 2026
8271025
Merge branch 'main' into dev/gokul/atan2-support
gokulkrishna98 Jul 6, 2026
8f15aca
Merge branch 'main' into dev/gokul/atan2-support
gokulkrishna98 Jul 7, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
106 changes: 106 additions & 0 deletions coreai_torch/_aten_to_core.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1547,6 +1547,111 @@ def replace_argmax(values_map: dict[str, Value], node: fx.Node, loc: Location) -
return result if keepdim else coreai.shrink_dims(result, [dim])


def replace_atan2(values_map: dict[str, Value], node: fx.Node, loc: Location) -> Value:
"""Lower atan2(y, x) using atan(y/x) with quadrant correction.

CoreAI has no native atan2, so it is decomposed as:
- 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.
atan2(0, 0) = 0 by convention.
"""
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)
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.
Comment thread
gokulkrishna98 marked this conversation as resolved.
y_is_zero = coreai.broadcasting_equal(y, zero)
x_is_zero = coreai.broadcasting_equal(x, zero)
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),
)

# ── 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)

# ── 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:
"""Converts aten.gather to coreai.gather_along_axis."""
x, index = _get_operands(values_map, node, [0, 2])
Expand Down Expand Up @@ -3465,6 +3570,7 @@ def sdpa_maskless(q: Value, k: Value, v: Value) -> Value:
"asin.default": replace_unary_ops,
"asinh.default": replace_unary_ops,
"atan.default": replace_unary_ops,
"atan2.default": replace_atan2,
"atanh.default": replace_unary_ops,
"_adaptive_avg_pool2d.default": replace_adaptive_avg_pool2d,
"_unsafe_view.default": replace_view,
Expand Down
74 changes: 74 additions & 0 deletions tests/ops/test_ops.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -608,6 +608,80 @@ def forward(self, x: Tensor) -> Tensor:
await validate_numerical_output(model=model, x=x, dynamic_shapes=dynamic_shapes)


class TestAtan2:
"""Tests for torch.atan2(y, x) — angle from the positive x-axis to the point (x, y)."""

class Atan2Model(nn.Module):
def forward(self, y: Tensor, x: Tensor) -> Tensor:
return torch.atan2(y, x)

@pytest.mark.parametrize("dynamic", [False, True])
@pytest.mark.parametrize(
"shape",
[
(4,),
(3, 4),
(2, 3, 4),
],
)
@pytest.mark.parametrize("dtype", [torch.float32, torch.float16])
async def test_basic(
self, shape: tuple[int, ...], dtype: torch.dtype, dynamic: bool
) -> None:
model = self.Atan2Model().eval()
y = torch.randn(shape, dtype=dtype)
x = torch.randn(shape, dtype=dtype)
dynamic_shapes = (
{"y": _all_dims_dynamic(y), "x": _all_dims_dynamic(x)} if dynamic else None
)
await validate_numerical_output(
model=model, y=y, x=x, dynamic_shapes=dynamic_shapes
)

async def test_x_zero(self) -> None:
"""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, 0.0])
x = torch.zeros(5)
await validate_numerical_output(model=model, y=y, x=x)

async def test_y_zero(self) -> None:
"""y = 0 with x > 0 → 0, x < 0 → π."""
model = self.Atan2Model().eval()
y = torch.zeros(4)
x = torch.tensor([1.0, -1.0, 2.0, -2.0])
await validate_numerical_output(model=model, y=y, x=x)

async def test_all_quadrants(self) -> None:
"""Cover all four quadrants and axes."""
model = self.Atan2Model().eval()
y = torch.tensor([1.0, 1.0, -1.0, -1.0, 0.0, 0.0, 1.0, -1.0])
x = torch.tensor([1.0, -1.0, 1.0, -1.0, 1.0, -1.0, 0.0, 0.0])
await validate_numerical_output(model=model, y=y, x=x)

async def test_broadcast_shapes(self) -> None:
model = self.Atan2Model().eval()
y = torch.randn(3, 4)
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",
[
Expand Down
133 changes: 133 additions & 0 deletions tests/ops/test_ops_ir.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1161,6 +1161,139 @@ def forward(self, x: Tensor) -> Tensor:
)


class TestAtan2IR:
def test_static(self) -> None:
class Atan2Model(nn.Module):
def forward(self, y: Tensor, x: Tensor) -> Tensor:
return torch.atan2(y, x)

ir = get_ir(Atan2Model().eval(), y=torch.rand(2, 3), x=torch.rand(2, 3))
filecheck_pattern(
ir,
check_file="""
// CHECK-LABEL: module {
// CHECK-NEXT: coreai.graph @main(%[[Y:.*]]: tensor<2x3xf32> {coreai.name = "y"}, %[[X:.*]]: tensor<2x3xf32> {coreai.name = "x"}) -> (tensor<2x3xf32> {coreai.name = "{{.*}}"}) attributes {__coreai_pure__} {
// CHECK: %[[NEG_INF:.*]] = coreai.constant dense<0xFF800000> : tensor<f32>
// CHECK: %[[POS_INF:.*]] = coreai.constant dense<0x7F800000> : tensor<f32>
// CHECK: %[[ZERO:.*]] = coreai.constant dense<0.000000e+00> : tensor<f32>
// CHECK: %[[ONE:.*]] = coreai.constant dense<1.000000e+00> : tensor<f32>
// CHECK: %[[PI:.*]] = coreai.constant dense<3.14159274> : tensor<f32>
// CHECK: %[[NEG_PI:.*]] = coreai.constant dense<-3.14159274> : tensor<f32>
// CHECK: %[[HPI:.*]] = coreai.constant dense<1.57079637> : tensor<f32>
// CHECK: %[[NHPI:.*]] = coreai.constant dense<-1.57079637> : tensor<f32>
// CHECK: %[[QPI:.*]] = coreai.constant dense<0.785398185> : tensor<f32>
// CHECK: %[[NQPI:.*]] = coreai.constant dense<-0.785398185> : tensor<f32>
// CHECK: %[[THREEQPI:.*]] = coreai.constant dense<2.3561945> : tensor<f32>
// CHECK: %[[NTHREEQPI:.*]] = coreai.constant dense<-2.3561945> : tensor<f32>
// CHECK: %[[Y_IS_ZERO:.*]] = coreai.decomposable.broadcasting_equal %[[Y]], %[[ZERO]]
// CHECK: %[[X_IS_ZERO:.*]] = coreai.decomposable.broadcasting_equal %[[X]], %[[ZERO]]
// CHECK: %[[Y_NEG_STRICT:.*]] = coreai.decomposable.broadcasting_greater %[[ZERO]], %[[Y]]
// CHECK: %[[RECIP_Y:.*]] = coreai.decomposable.broadcasting_divide %[[ONE]], %[[Y]]
// CHECK: %[[RECIP_Y_NEG:.*]] = coreai.decomposable.broadcasting_greater %[[ZERO]], %[[RECIP_Y]]
// CHECK: %[[Y_ZERO_NEG:.*]] = coreai.decomposable.broadcasting_and %[[Y_IS_ZERO]], %[[RECIP_Y_NEG]]
// CHECK: %[[Y_NEG:.*]] = coreai.decomposable.broadcasting_or %[[Y_NEG_STRICT]], %[[Y_ZERO_NEG]]
// CHECK: %[[X_NEG_STRICT:.*]] = coreai.decomposable.broadcasting_greater %[[ZERO]], %[[X]]
// CHECK: %[[RECIP_X:.*]] = coreai.decomposable.broadcasting_divide %[[ONE]], %[[X]]
// CHECK: %[[RECIP_X_NEG:.*]] = coreai.decomposable.broadcasting_greater %[[ZERO]], %[[RECIP_X]]
// CHECK: %[[X_ZERO_NEG:.*]] = coreai.decomposable.broadcasting_and %[[X_IS_ZERO]], %[[RECIP_X_NEG]]
// CHECK: %[[X_NEG:.*]] = coreai.decomposable.broadcasting_or %[[X_NEG_STRICT]], %[[X_ZERO_NEG]]
// CHECK: %[[X_IS_POS_INF:.*]] = coreai.decomposable.broadcasting_equal %[[X]], %[[POS_INF]]
// CHECK: %[[X_IS_NEG_INF:.*]] = coreai.decomposable.broadcasting_equal %[[X]], %[[NEG_INF]]
// CHECK: %[[X_IS_INF:.*]] = coreai.decomposable.broadcasting_or %[[X_IS_POS_INF]], %[[X_IS_NEG_INF]]
// CHECK: %[[Y_IS_POS_INF:.*]] = coreai.decomposable.broadcasting_equal %[[Y]], %[[POS_INF]]
// CHECK: %[[Y_IS_NEG_INF:.*]] = coreai.decomposable.broadcasting_equal %[[Y]], %[[NEG_INF]]
// CHECK: %[[Y_IS_INF:.*]] = coreai.decomposable.broadcasting_or %[[Y_IS_POS_INF]], %[[Y_IS_NEG_INF]]
// CHECK: %[[BOTH_INF:.*]] = coreai.decomposable.broadcasting_and %[[X_IS_INF]], %[[Y_IS_INF]]
// CHECK: %[[BASE:.*]] = coreai.atan
// CHECK: %[[RESULT:.*]] = coreai.decomposable.broadcasting_where %[[BOTH_INF]],
// CHECK-NEXT: coreai.output %[[RESULT]] : tensor<2x3xf32>
// CHECK-NEXT: }
// CHECK-NEXT: }
""",
)

def test_dynamic(self) -> None:
class Atan2Model(nn.Module):
def forward(self, y: Tensor, x: Tensor) -> Tensor:
return torch.atan2(y, x)

y = torch.rand(2, 3)
x = torch.rand(2, 3)
ir = get_ir(
Atan2Model().eval(),
y=y,
x=x,
dynamic_shapes={"y": _all_dims_dynamic(y), "x": _all_dims_dynamic(x)},
)
filecheck_pattern(
ir,
check_file="""
// CHECK-LABEL: module {
// 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__} {
// CHECK: %[[NEG_INF:.*]] = coreai.constant dense<0xFF800000> : tensor<f32>
// CHECK: %[[POS_INF:.*]] = coreai.constant dense<0x7F800000> : tensor<f32>
// CHECK: %[[ZERO:.*]] = coreai.constant dense<0.000000e+00> : tensor<f32>
// CHECK: %[[ONE:.*]] = coreai.constant dense<1.000000e+00> : tensor<f32>
// CHECK: %[[PI:.*]] = coreai.constant dense<3.14159274> : tensor<f32>
// CHECK: %[[NEG_PI:.*]] = coreai.constant dense<-3.14159274> : tensor<f32>
// CHECK: %[[HPI:.*]] = coreai.constant dense<1.57079637> : tensor<f32>
// CHECK: %[[NHPI:.*]] = coreai.constant dense<-1.57079637> : tensor<f32>
// CHECK: %[[QPI:.*]] = coreai.constant dense<0.785398185> : tensor<f32>
// CHECK: %[[NQPI:.*]] = coreai.constant dense<-0.785398185> : tensor<f32>
// CHECK: %[[THREEQPI:.*]] = coreai.constant dense<2.3561945> : tensor<f32>
// CHECK: %[[NTHREEQPI:.*]] = coreai.constant dense<-2.3561945> : tensor<f32>
// CHECK: %[[Y_IS_ZERO:.*]] = coreai.decomposable.broadcasting_equal %[[Y]], %[[ZERO]]
// CHECK: %[[X_IS_ZERO:.*]] = coreai.decomposable.broadcasting_equal %[[X]], %[[ZERO]]
// CHECK: %[[Y_NEG_STRICT:.*]] = coreai.decomposable.broadcasting_greater %[[ZERO]], %[[Y]]
// CHECK: %[[RECIP_Y:.*]] = coreai.decomposable.broadcasting_divide %[[ONE]], %[[Y]]
// CHECK: %[[RECIP_Y_NEG:.*]] = coreai.decomposable.broadcasting_greater %[[ZERO]], %[[RECIP_Y]]
// CHECK: %[[Y_ZERO_NEG:.*]] = coreai.decomposable.broadcasting_and %[[Y_IS_ZERO]], %[[RECIP_Y_NEG]]
// CHECK: %[[Y_NEG:.*]] = coreai.decomposable.broadcasting_or %[[Y_NEG_STRICT]], %[[Y_ZERO_NEG]]
// CHECK: %[[X_NEG_STRICT:.*]] = coreai.decomposable.broadcasting_greater %[[ZERO]], %[[X]]
// CHECK: %[[RECIP_X:.*]] = coreai.decomposable.broadcasting_divide %[[ONE]], %[[X]]
// CHECK: %[[RECIP_X_NEG:.*]] = coreai.decomposable.broadcasting_greater %[[ZERO]], %[[RECIP_X]]
// CHECK: %[[X_ZERO_NEG:.*]] = coreai.decomposable.broadcasting_and %[[X_IS_ZERO]], %[[RECIP_X_NEG]]
// CHECK: %[[X_NEG:.*]] = coreai.decomposable.broadcasting_or %[[X_NEG_STRICT]], %[[X_ZERO_NEG]]
// CHECK: %[[X_IS_POS_INF:.*]] = coreai.decomposable.broadcasting_equal %[[X]], %[[POS_INF]]
// CHECK: %[[X_IS_NEG_INF:.*]] = coreai.decomposable.broadcasting_equal %[[X]], %[[NEG_INF]]
// CHECK: %[[X_IS_INF:.*]] = coreai.decomposable.broadcasting_or %[[X_IS_POS_INF]], %[[X_IS_NEG_INF]]
// CHECK: %[[Y_IS_POS_INF:.*]] = coreai.decomposable.broadcasting_equal %[[Y]], %[[POS_INF]]
// CHECK: %[[Y_IS_NEG_INF:.*]] = coreai.decomposable.broadcasting_equal %[[Y]], %[[NEG_INF]]
// CHECK: %[[Y_IS_INF:.*]] = coreai.decomposable.broadcasting_or %[[Y_IS_POS_INF]], %[[Y_IS_NEG_INF]]
// CHECK: %[[BOTH_INF:.*]] = coreai.decomposable.broadcasting_and %[[X_IS_INF]], %[[Y_IS_INF]]
// CHECK: %[[BASE:.*]] = coreai.atan
// CHECK: %[[RESULT:.*]] = coreai.decomposable.broadcasting_where %[[BOTH_INF]],
// CHECK-NEXT: coreai.output %[[RESULT]] : tensor<?x?xf32>
// CHECK-NEXT: }
// CHECK-NEXT: }
""",
)

def test_1d(self) -> None:
class Atan2Model(nn.Module):
def forward(self, y: Tensor, x: Tensor) -> Tensor:
return torch.atan2(y, x)

ir = get_ir(Atan2Model().eval(), y=torch.rand(4), x=torch.rand(4))
filecheck_pattern(
ir,
check_file="""
// CHECK-LABEL: module {
// CHECK-NEXT: coreai.graph @main(%[[Y:.*]]: tensor<4xf32> {coreai.name = "y"}, %[[X:.*]]: tensor<4xf32> {coreai.name = "x"}) -> (tensor<4xf32> {coreai.name = "{{.*}}"}) attributes {__coreai_pure__} {
// CHECK: %[[ZERO:.*]] = coreai.constant dense<0.000000e+00> : tensor<f32>
// CHECK: %[[ONE:.*]] = coreai.constant dense<1.000000e+00> : tensor<f32>
// CHECK: %[[Y_NEG:.*]] = coreai.decomposable.broadcasting_or
// CHECK: %[[X_NEG:.*]] = coreai.decomposable.broadcasting_or
// CHECK: %[[BOTH_INF:.*]] = coreai.decomposable.broadcasting_and
// CHECK: %[[BASE:.*]] = coreai.atan
// CHECK: %[[RESULT:.*]] = coreai.decomposable.broadcasting_where %[[BOTH_INF]],
// CHECK-NEXT: coreai.output %[[RESULT]] : tensor<4xf32>
// CHECK-NEXT: }
// CHECK-NEXT: }
""",
)


class TestAvgPool2dIR:
def test_static(self) -> None:
class AvgPool2dModel(nn.Module):
Expand Down