Skip to content

Commit a43cc84

Browse files
gokulkrishna98gokulkrishna98
authored
[converter] implement aten::atan2 conversion (#23)
## Description Adds a converter for aten.atan2.default. coreai has no native atan2 op, so the conversion decomposes it into atan(y/x) with per-quadrant correction: - x > 0: atan(y/x) - x < 0, y ≥ 0: atan(y/x) + π - x < 0, y < 0: atan(y/x) − π - x = 0, y > 0: π/2 - x = 0, y < 0: −π/2 - x = 0, y = 0: 0 Division by zero when x=0 is guarded with broadcasting_where before the divide, then the x=0 result is selected in a final where at the end. ## Testing: - Adds numerical tests (shapes 1D/2D/3D, float32/float16, static/dynamic, axis-aligned edge cases) and IR FileCheck tests (static, dynamic, 1D). - python unit tests.
1 parent a374b48 commit a43cc84

3 files changed

Lines changed: 313 additions & 0 deletions

File tree

‎coreai_torch/_aten_to_core.py‎

Lines changed: 106 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1547,6 +1547,111 @@ def replace_argmax(values_map: dict[str, Value], node: fx.Node, loc: Location) -
15471547
return result if keepdim else coreai.shrink_dims(result, [dim])
15481548

15491549

1550+
def replace_atan2(values_map: dict[str, Value], node: fx.Node, loc: Location) -> Value:
1551+
"""Lower atan2(y, x) using atan(y/x) with quadrant correction.
1552+
1553+
CoreAI has no native atan2, so it is decomposed as:
1554+
- x != 0, finite: atan(y/x) adjusted by ±π for the correct quadrant.
1555+
- x == +0: ±π/2 for non-zero y, 0 for y = 0.
1556+
- x == -0: ±π for all y (including ±0 → ±π per IEEE-754).
1557+
- both infinite: ±π/4 or ±3π/4 per IEEE-754.
1558+
1559+
Signed-zero handling: IEEE-754 treats -0.0 as distinct from +0.0 for atan2
1560+
(e.g. atan2(-0, -1) = -π, not +π). The 1/v trick — 1/-0.0 = -inf — is used
1561+
to detect the sign bit of zero inputs so that y_neg and x_neg are correct
1562+
for -0.0 inputs without misclassifying ±inf (which use the strict > path).
1563+
1564+
When x=0, x is replaced with 1 before the divide solely to avoid NaN/inf; that
1565+
intermediate result is discarded by the final where-select.
1566+
atan2(0, 0) = 0 by convention.
1567+
"""
1568+
y, x = _get_operands(values_map, node, [0, 1])
1569+
ele_type = x.type.element_type
1570+
1571+
zero = coreai.constant(0.0, dtype=ele_type)
1572+
one = coreai.constant(1.0, dtype=ele_type)
1573+
pi = coreai.constant(np.pi, dtype=ele_type)
1574+
neg_pi = coreai.constant(-np.pi, dtype=ele_type)
1575+
half_pi = coreai.constant(np.pi / 2.0, dtype=ele_type)
1576+
neg_half_pi = coreai.constant(-np.pi / 2.0, dtype=ele_type)
1577+
quarter_pi = coreai.constant(np.pi / 4.0, dtype=ele_type)
1578+
neg_quarter_pi = coreai.constant(-np.pi / 4.0, dtype=ele_type)
1579+
three_quarter_pi = coreai.constant(3.0 * np.pi / 4.0, dtype=ele_type)
1580+
neg_three_quarter_pi = coreai.constant(-3.0 * np.pi / 4.0, dtype=ele_type)
1581+
1582+
# ── signed-zero-aware sign predicates ─────────────────────────────────────
1583+
# 1 / -0.0 = -inf (IEEE-754), so (0 > 1/v) is True iff v = -0.0. Combine with
1584+
# the strict > predicate (handles ±inf and non-zero finites) via OR.
1585+
y_is_zero = coreai.broadcasting_equal(y, zero)
1586+
x_is_zero = coreai.broadcasting_equal(x, zero)
1587+
y_neg = coreai.broadcasting_or(
1588+
coreai.broadcasting_greater(zero, y),
1589+
coreai.broadcasting_and(
1590+
y_is_zero,
1591+
coreai.broadcasting_greater(zero, coreai.broadcasting_divide(one, y)),
1592+
),
1593+
)
1594+
x_neg = coreai.broadcasting_or(
1595+
coreai.broadcasting_greater(zero, x),
1596+
coreai.broadcasting_and(
1597+
x_is_zero,
1598+
coreai.broadcasting_greater(zero, coreai.broadcasting_divide(one, x)),
1599+
),
1600+
)
1601+
x_is_neg_zero = coreai.broadcasting_and(
1602+
x_is_zero,
1603+
coreai.broadcasting_greater(zero, coreai.broadcasting_divide(one, x)),
1604+
)
1605+
1606+
# ── both-infinite branch ──────────────────────────────────────────────────
1607+
# atan(inf/inf) = atan(NaN) = NaN; handle before the divide.
1608+
pos_inf = coreai.constant(float("inf"), dtype=ele_type)
1609+
neg_inf = coreai.constant(float("-inf"), dtype=ele_type)
1610+
x_is_inf = coreai.broadcasting_or(
1611+
coreai.broadcasting_equal(x, pos_inf), coreai.broadcasting_equal(x, neg_inf)
1612+
)
1613+
y_is_inf = coreai.broadcasting_or(
1614+
coreai.broadcasting_equal(y, pos_inf), coreai.broadcasting_equal(y, neg_inf)
1615+
)
1616+
both_inf = coreai.broadcasting_and(x_is_inf, y_is_inf)
1617+
inf_result = coreai.broadcasting_where(
1618+
y_neg,
1619+
coreai.broadcasting_where(x_neg, neg_three_quarter_pi, neg_quarter_pi),
1620+
coreai.broadcasting_where(x_neg, three_quarter_pi, quarter_pi),
1621+
)
1622+
1623+
# ── x = 0 branch ──────────────────────────────────────────────────────────
1624+
# x = +0: ±π/2 for strictly ±y, 0 when y = 0.
1625+
# x = -0: ±π for all y (y_neg covers y = -0.0 via the 1/y trick above).
1626+
y_pos_strict = coreai.broadcasting_greater(y, zero)
1627+
y_neg_strict = coreai.broadcasting_greater(zero, y)
1628+
pos_x_zero_result = coreai.broadcasting_where(
1629+
y_pos_strict,
1630+
half_pi,
1631+
coreai.broadcasting_where(y_neg_strict, neg_half_pi, zero),
1632+
)
1633+
neg_x_zero_result = coreai.broadcasting_where(y_neg, neg_pi, pi)
1634+
zero_result = coreai.broadcasting_where(
1635+
x_is_neg_zero, neg_x_zero_result, pos_x_zero_result
1636+
)
1637+
1638+
# ── finite nonzero x branch ────────────────────────────────────────────────
1639+
# Avoid division by zero: substitute x = 1 when x = 0; result discarded by
1640+
# the outer where-select.
1641+
x_safe = coreai.broadcasting_where(x_is_zero, one, x)
1642+
base = coreai.atan(coreai.broadcasting_divide(y, x_safe))
1643+
correction = coreai.broadcasting_where(
1644+
y_neg,
1645+
coreai.broadcasting_sub(base, pi),
1646+
coreai.broadcasting_add(base, pi),
1647+
)
1648+
nonzero_result = coreai.broadcasting_where(x_neg, correction, base)
1649+
1650+
# ── combine ────────────────────────────────────────────────────────────────
1651+
result = coreai.broadcasting_where(x_is_zero, zero_result, nonzero_result)
1652+
return coreai.broadcasting_where(both_inf, inf_result, result)
1653+
1654+
15501655
def replace_gather(values_map: dict[str, Value], node: fx.Node, loc: Location) -> Value:
15511656
"""Converts aten.gather to coreai.gather_along_axis."""
15521657
x, index = _get_operands(values_map, node, [0, 2])
@@ -3465,6 +3570,7 @@ def sdpa_maskless(q: Value, k: Value, v: Value) -> Value:
34653570
"asin.default": replace_unary_ops,
34663571
"asinh.default": replace_unary_ops,
34673572
"atan.default": replace_unary_ops,
3573+
"atan2.default": replace_atan2,
34683574
"atanh.default": replace_unary_ops,
34693575
"_adaptive_avg_pool2d.default": replace_adaptive_avg_pool2d,
34703576
"_unsafe_view.default": replace_view,

‎tests/ops/test_ops.py‎

Lines changed: 74 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -608,6 +608,80 @@ def forward(self, x: Tensor) -> Tensor:
608608
await validate_numerical_output(model=model, x=x, dynamic_shapes=dynamic_shapes)
609609

610610

611+
class TestAtan2:
612+
"""Tests for torch.atan2(y, x) — angle from the positive x-axis to the point (x, y)."""
613+
614+
class Atan2Model(nn.Module):
615+
def forward(self, y: Tensor, x: Tensor) -> Tensor:
616+
return torch.atan2(y, x)
617+
618+
@pytest.mark.parametrize("dynamic", [False, True])
619+
@pytest.mark.parametrize(
620+
"shape",
621+
[
622+
(4,),
623+
(3, 4),
624+
(2, 3, 4),
625+
],
626+
)
627+
@pytest.mark.parametrize("dtype", [torch.float32, torch.float16])
628+
async def test_basic(
629+
self, shape: tuple[int, ...], dtype: torch.dtype, dynamic: bool
630+
) -> None:
631+
model = self.Atan2Model().eval()
632+
y = torch.randn(shape, dtype=dtype)
633+
x = torch.randn(shape, dtype=dtype)
634+
dynamic_shapes = (
635+
{"y": _all_dims_dynamic(y), "x": _all_dims_dynamic(x)} if dynamic else None
636+
)
637+
await validate_numerical_output(
638+
model=model, y=y, x=x, dynamic_shapes=dynamic_shapes
639+
)
640+
641+
async def test_x_zero(self) -> None:
642+
"""x = 0 should yield ±π/2 depending on sign of y; (0, 0) → 0 by convention."""
643+
model = self.Atan2Model().eval()
644+
y = torch.tensor([1.0, -1.0, 2.0, -2.0, 0.0])
645+
x = torch.zeros(5)
646+
await validate_numerical_output(model=model, y=y, x=x)
647+
648+
async def test_y_zero(self) -> None:
649+
"""y = 0 with x > 0 → 0, x < 0 → π."""
650+
model = self.Atan2Model().eval()
651+
y = torch.zeros(4)
652+
x = torch.tensor([1.0, -1.0, 2.0, -2.0])
653+
await validate_numerical_output(model=model, y=y, x=x)
654+
655+
async def test_all_quadrants(self) -> None:
656+
"""Cover all four quadrants and axes."""
657+
model = self.Atan2Model().eval()
658+
y = torch.tensor([1.0, 1.0, -1.0, -1.0, 0.0, 0.0, 1.0, -1.0])
659+
x = torch.tensor([1.0, -1.0, 1.0, -1.0, 1.0, -1.0, 0.0, 0.0])
660+
await validate_numerical_output(model=model, y=y, x=x)
661+
662+
async def test_broadcast_shapes(self) -> None:
663+
model = self.Atan2Model().eval()
664+
y = torch.randn(3, 4)
665+
x = torch.randn(4)
666+
await validate_numerical_output(model=model, y=y, x=x)
667+
668+
async def test_signed_zeros(self) -> None:
669+
"""IEEE-754 signed-zero cases: atan2(-0, x) and atan2(y, -0)."""
670+
model = self.Atan2Model().eval()
671+
# y = -0.0 with various x signs
672+
y = torch.tensor([-0.0, -0.0, -0.0, 0.0])
673+
x = torch.tensor([-1.0, 1.0, -0.0, -0.0])
674+
await validate_numerical_output(model=model, y=y, x=x)
675+
676+
async def test_infinities(self) -> None:
677+
"""IEEE-754 both-infinite cases: atan2(±inf, ±inf) → ±π/4 or ±3π/4."""
678+
model = self.Atan2Model().eval()
679+
inf = float("inf")
680+
y = torch.tensor([inf, inf, -inf, -inf])
681+
x = torch.tensor([inf, -inf, inf, -inf])
682+
await validate_numerical_output(model=model, y=y, x=x)
683+
684+
611685
@pytest.mark.parametrize(
612686
"x",
613687
[

‎tests/ops/test_ops_ir.py‎

Lines changed: 133 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1161,6 +1161,139 @@ def forward(self, x: Tensor) -> Tensor:
11611161
)
11621162

11631163

1164+
class TestAtan2IR:
1165+
def test_static(self) -> None:
1166+
class Atan2Model(nn.Module):
1167+
def forward(self, y: Tensor, x: Tensor) -> Tensor:
1168+
return torch.atan2(y, x)
1169+
1170+
ir = get_ir(Atan2Model().eval(), y=torch.rand(2, 3), x=torch.rand(2, 3))
1171+
filecheck_pattern(
1172+
ir,
1173+
check_file="""
1174+
// CHECK-LABEL: module {
1175+
// CHECK-NEXT: coreai.graph @main(%[[Y:.*]]: tensor<2x3xf32> {coreai.name = "y"}, %[[X:.*]]: tensor<2x3xf32> {coreai.name = "x"}) -> (tensor<2x3xf32> {coreai.name = "{{.*}}"}) attributes {__coreai_pure__} {
1176+
// CHECK: %[[NEG_INF:.*]] = coreai.constant dense<0xFF800000> : tensor<f32>
1177+
// CHECK: %[[POS_INF:.*]] = coreai.constant dense<0x7F800000> : tensor<f32>
1178+
// CHECK: %[[ZERO:.*]] = coreai.constant dense<0.000000e+00> : tensor<f32>
1179+
// CHECK: %[[ONE:.*]] = coreai.constant dense<1.000000e+00> : tensor<f32>
1180+
// CHECK: %[[PI:.*]] = coreai.constant dense<3.14159274> : tensor<f32>
1181+
// CHECK: %[[NEG_PI:.*]] = coreai.constant dense<-3.14159274> : tensor<f32>
1182+
// CHECK: %[[HPI:.*]] = coreai.constant dense<1.57079637> : tensor<f32>
1183+
// CHECK: %[[NHPI:.*]] = coreai.constant dense<-1.57079637> : tensor<f32>
1184+
// CHECK: %[[QPI:.*]] = coreai.constant dense<0.785398185> : tensor<f32>
1185+
// CHECK: %[[NQPI:.*]] = coreai.constant dense<-0.785398185> : tensor<f32>
1186+
// CHECK: %[[THREEQPI:.*]] = coreai.constant dense<2.3561945> : tensor<f32>
1187+
// CHECK: %[[NTHREEQPI:.*]] = coreai.constant dense<-2.3561945> : tensor<f32>
1188+
// CHECK: %[[Y_IS_ZERO:.*]] = coreai.decomposable.broadcasting_equal %[[Y]], %[[ZERO]]
1189+
// CHECK: %[[X_IS_ZERO:.*]] = coreai.decomposable.broadcasting_equal %[[X]], %[[ZERO]]
1190+
// CHECK: %[[Y_NEG_STRICT:.*]] = coreai.decomposable.broadcasting_greater %[[ZERO]], %[[Y]]
1191+
// CHECK: %[[RECIP_Y:.*]] = coreai.decomposable.broadcasting_divide %[[ONE]], %[[Y]]
1192+
// CHECK: %[[RECIP_Y_NEG:.*]] = coreai.decomposable.broadcasting_greater %[[ZERO]], %[[RECIP_Y]]
1193+
// CHECK: %[[Y_ZERO_NEG:.*]] = coreai.decomposable.broadcasting_and %[[Y_IS_ZERO]], %[[RECIP_Y_NEG]]
1194+
// CHECK: %[[Y_NEG:.*]] = coreai.decomposable.broadcasting_or %[[Y_NEG_STRICT]], %[[Y_ZERO_NEG]]
1195+
// CHECK: %[[X_NEG_STRICT:.*]] = coreai.decomposable.broadcasting_greater %[[ZERO]], %[[X]]
1196+
// CHECK: %[[RECIP_X:.*]] = coreai.decomposable.broadcasting_divide %[[ONE]], %[[X]]
1197+
// CHECK: %[[RECIP_X_NEG:.*]] = coreai.decomposable.broadcasting_greater %[[ZERO]], %[[RECIP_X]]
1198+
// CHECK: %[[X_ZERO_NEG:.*]] = coreai.decomposable.broadcasting_and %[[X_IS_ZERO]], %[[RECIP_X_NEG]]
1199+
// CHECK: %[[X_NEG:.*]] = coreai.decomposable.broadcasting_or %[[X_NEG_STRICT]], %[[X_ZERO_NEG]]
1200+
// CHECK: %[[X_IS_POS_INF:.*]] = coreai.decomposable.broadcasting_equal %[[X]], %[[POS_INF]]
1201+
// CHECK: %[[X_IS_NEG_INF:.*]] = coreai.decomposable.broadcasting_equal %[[X]], %[[NEG_INF]]
1202+
// CHECK: %[[X_IS_INF:.*]] = coreai.decomposable.broadcasting_or %[[X_IS_POS_INF]], %[[X_IS_NEG_INF]]
1203+
// CHECK: %[[Y_IS_POS_INF:.*]] = coreai.decomposable.broadcasting_equal %[[Y]], %[[POS_INF]]
1204+
// CHECK: %[[Y_IS_NEG_INF:.*]] = coreai.decomposable.broadcasting_equal %[[Y]], %[[NEG_INF]]
1205+
// CHECK: %[[Y_IS_INF:.*]] = coreai.decomposable.broadcasting_or %[[Y_IS_POS_INF]], %[[Y_IS_NEG_INF]]
1206+
// CHECK: %[[BOTH_INF:.*]] = coreai.decomposable.broadcasting_and %[[X_IS_INF]], %[[Y_IS_INF]]
1207+
// CHECK: %[[BASE:.*]] = coreai.atan
1208+
// CHECK: %[[RESULT:.*]] = coreai.decomposable.broadcasting_where %[[BOTH_INF]],
1209+
// CHECK-NEXT: coreai.output %[[RESULT]] : tensor<2x3xf32>
1210+
// CHECK-NEXT: }
1211+
// CHECK-NEXT: }
1212+
""",
1213+
)
1214+
1215+
def test_dynamic(self) -> None:
1216+
class Atan2Model(nn.Module):
1217+
def forward(self, y: Tensor, x: Tensor) -> Tensor:
1218+
return torch.atan2(y, x)
1219+
1220+
y = torch.rand(2, 3)
1221+
x = torch.rand(2, 3)
1222+
ir = get_ir(
1223+
Atan2Model().eval(),
1224+
y=y,
1225+
x=x,
1226+
dynamic_shapes={"y": _all_dims_dynamic(y), "x": _all_dims_dynamic(x)},
1227+
)
1228+
filecheck_pattern(
1229+
ir,
1230+
check_file="""
1231+
// CHECK-LABEL: module {
1232+
// 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__} {
1233+
// CHECK: %[[NEG_INF:.*]] = coreai.constant dense<0xFF800000> : tensor<f32>
1234+
// CHECK: %[[POS_INF:.*]] = coreai.constant dense<0x7F800000> : tensor<f32>
1235+
// CHECK: %[[ZERO:.*]] = coreai.constant dense<0.000000e+00> : tensor<f32>
1236+
// CHECK: %[[ONE:.*]] = coreai.constant dense<1.000000e+00> : tensor<f32>
1237+
// CHECK: %[[PI:.*]] = coreai.constant dense<3.14159274> : tensor<f32>
1238+
// CHECK: %[[NEG_PI:.*]] = coreai.constant dense<-3.14159274> : tensor<f32>
1239+
// CHECK: %[[HPI:.*]] = coreai.constant dense<1.57079637> : tensor<f32>
1240+
// CHECK: %[[NHPI:.*]] = coreai.constant dense<-1.57079637> : tensor<f32>
1241+
// CHECK: %[[QPI:.*]] = coreai.constant dense<0.785398185> : tensor<f32>
1242+
// CHECK: %[[NQPI:.*]] = coreai.constant dense<-0.785398185> : tensor<f32>
1243+
// CHECK: %[[THREEQPI:.*]] = coreai.constant dense<2.3561945> : tensor<f32>
1244+
// CHECK: %[[NTHREEQPI:.*]] = coreai.constant dense<-2.3561945> : tensor<f32>
1245+
// CHECK: %[[Y_IS_ZERO:.*]] = coreai.decomposable.broadcasting_equal %[[Y]], %[[ZERO]]
1246+
// CHECK: %[[X_IS_ZERO:.*]] = coreai.decomposable.broadcasting_equal %[[X]], %[[ZERO]]
1247+
// CHECK: %[[Y_NEG_STRICT:.*]] = coreai.decomposable.broadcasting_greater %[[ZERO]], %[[Y]]
1248+
// CHECK: %[[RECIP_Y:.*]] = coreai.decomposable.broadcasting_divide %[[ONE]], %[[Y]]
1249+
// CHECK: %[[RECIP_Y_NEG:.*]] = coreai.decomposable.broadcasting_greater %[[ZERO]], %[[RECIP_Y]]
1250+
// CHECK: %[[Y_ZERO_NEG:.*]] = coreai.decomposable.broadcasting_and %[[Y_IS_ZERO]], %[[RECIP_Y_NEG]]
1251+
// CHECK: %[[Y_NEG:.*]] = coreai.decomposable.broadcasting_or %[[Y_NEG_STRICT]], %[[Y_ZERO_NEG]]
1252+
// CHECK: %[[X_NEG_STRICT:.*]] = coreai.decomposable.broadcasting_greater %[[ZERO]], %[[X]]
1253+
// CHECK: %[[RECIP_X:.*]] = coreai.decomposable.broadcasting_divide %[[ONE]], %[[X]]
1254+
// CHECK: %[[RECIP_X_NEG:.*]] = coreai.decomposable.broadcasting_greater %[[ZERO]], %[[RECIP_X]]
1255+
// CHECK: %[[X_ZERO_NEG:.*]] = coreai.decomposable.broadcasting_and %[[X_IS_ZERO]], %[[RECIP_X_NEG]]
1256+
// CHECK: %[[X_NEG:.*]] = coreai.decomposable.broadcasting_or %[[X_NEG_STRICT]], %[[X_ZERO_NEG]]
1257+
// CHECK: %[[X_IS_POS_INF:.*]] = coreai.decomposable.broadcasting_equal %[[X]], %[[POS_INF]]
1258+
// CHECK: %[[X_IS_NEG_INF:.*]] = coreai.decomposable.broadcasting_equal %[[X]], %[[NEG_INF]]
1259+
// CHECK: %[[X_IS_INF:.*]] = coreai.decomposable.broadcasting_or %[[X_IS_POS_INF]], %[[X_IS_NEG_INF]]
1260+
// CHECK: %[[Y_IS_POS_INF:.*]] = coreai.decomposable.broadcasting_equal %[[Y]], %[[POS_INF]]
1261+
// CHECK: %[[Y_IS_NEG_INF:.*]] = coreai.decomposable.broadcasting_equal %[[Y]], %[[NEG_INF]]
1262+
// CHECK: %[[Y_IS_INF:.*]] = coreai.decomposable.broadcasting_or %[[Y_IS_POS_INF]], %[[Y_IS_NEG_INF]]
1263+
// CHECK: %[[BOTH_INF:.*]] = coreai.decomposable.broadcasting_and %[[X_IS_INF]], %[[Y_IS_INF]]
1264+
// CHECK: %[[BASE:.*]] = coreai.atan
1265+
// CHECK: %[[RESULT:.*]] = coreai.decomposable.broadcasting_where %[[BOTH_INF]],
1266+
// CHECK-NEXT: coreai.output %[[RESULT]] : tensor<?x?xf32>
1267+
// CHECK-NEXT: }
1268+
// CHECK-NEXT: }
1269+
""",
1270+
)
1271+
1272+
def test_1d(self) -> None:
1273+
class Atan2Model(nn.Module):
1274+
def forward(self, y: Tensor, x: Tensor) -> Tensor:
1275+
return torch.atan2(y, x)
1276+
1277+
ir = get_ir(Atan2Model().eval(), y=torch.rand(4), x=torch.rand(4))
1278+
filecheck_pattern(
1279+
ir,
1280+
check_file="""
1281+
// CHECK-LABEL: module {
1282+
// CHECK-NEXT: coreai.graph @main(%[[Y:.*]]: tensor<4xf32> {coreai.name = "y"}, %[[X:.*]]: tensor<4xf32> {coreai.name = "x"}) -> (tensor<4xf32> {coreai.name = "{{.*}}"}) attributes {__coreai_pure__} {
1283+
// CHECK: %[[ZERO:.*]] = coreai.constant dense<0.000000e+00> : tensor<f32>
1284+
// CHECK: %[[ONE:.*]] = coreai.constant dense<1.000000e+00> : tensor<f32>
1285+
// CHECK: %[[Y_NEG:.*]] = coreai.decomposable.broadcasting_or
1286+
// CHECK: %[[X_NEG:.*]] = coreai.decomposable.broadcasting_or
1287+
// CHECK: %[[BOTH_INF:.*]] = coreai.decomposable.broadcasting_and
1288+
// CHECK: %[[BASE:.*]] = coreai.atan
1289+
// CHECK: %[[RESULT:.*]] = coreai.decomposable.broadcasting_where %[[BOTH_INF]],
1290+
// CHECK-NEXT: coreai.output %[[RESULT]] : tensor<4xf32>
1291+
// CHECK-NEXT: }
1292+
// CHECK-NEXT: }
1293+
""",
1294+
)
1295+
1296+
11641297
class TestAvgPool2dIR:
11651298
def test_static(self) -> None:
11661299
class AvgPool2dModel(nn.Module):

0 commit comments

Comments
 (0)