Merge branch 'main' into fix/dequantize-negative-axis

Author: jakesabathia2

.github/workflows/ci.yml

@@ -34,4 +34,4 @@ jobs:
         run: |
           command -v uv >/dev/null 2>&1 || curl -LsSf https://astral.sh/uv/install.sh | sh
           echo "$HOME/.local/bin" >> "$GITHUB_PATH"
-      - run: uv run --extra test pytest tests/ -n auto -m "not slow"
+      - run: uv run --extra test pytest tests/ -n auto -m "not slow and not dsl"

coreai_torch/__version__.py

@@ -5,4 +5,4 @@
 
 """Version information for coreai-torch."""
 
-__version__ = "0.4.0"
+__version__ = "0.4.1"

coreai_torch/_aten_to_core.py

@@ -53,6 +53,9 @@
 from ._utils import (
     get_operands as _get_operands,
 )
+from ._utils import (
+    replace_pad_with_mode as _replace_pad_with_mode,
+)
 
 INT32_MAX: int = 2147483647
 
@@ -644,20 +647,28 @@ def replace_arange_start_step(
         else coreai.constant(1, dtype=start.type.element_type)
     )
 
-    # coreai.range_ requires scalar (rank-0) operands that share an element
-    # type. aten.arange promotes mixed-type scalars internally; we replicate
-    # that here by squeezing each operand to rank-0 and casting to the FX
-    # node's output dtype before the op.
+    # When ALL operands are integer-typed, keep them as si32 so coreai.range_
+    # can infer a static output shape, then cast the result to the requested
+    # dtype. If any operand is float (e.g. arange(0, 5, 0.5)), fall back to
+    # target_type — truncating a float step to si32 would corrupt the values.
     target_type = get_output_element_type_from_node(node)
+    si32 = IntegerType.get_signed(32)
+    all_integer = all(
+        isinstance(v.type.element_type, IntegerType) for v in (start, end, step)
+    )
+    range_type = si32 if all_integer else target_type
 
     def to_scalar(v: Value) -> Value:
         if v.type.rank > 0:
             v = coreai.shrink_dims(v, list(range(v.type.rank)))
-        if v.type.element_type != target_type:
-            v = coreai.cast(v, target_type)
+        if v.type.element_type != range_type:
+            v = coreai.cast(v, range_type)
         return v
 
-    return coreai.range_(to_scalar(start), to_scalar(end), to_scalar(step))
+    result = coreai.range_(to_scalar(start), to_scalar(end), to_scalar(step))
+    if result.type.element_type != target_type:
+        result = coreai.cast(result, target_type)
+    return result
 
 
 def replace_batch_norm(
@@ -1109,6 +1120,20 @@ def replace_constant_pad_nd(
     return result
 
 
+def replace_reflection_pad(
+    values_map: dict[str, Value], node: fx.Node, loc: Location
+) -> Value:
+    """aten.reflection_pad{1,2,3}d.default -> coreai.pad<reflect>."""
+    return _replace_pad_with_mode(values_map, node, loc, "reflect")
+
+
+def replace_replication_pad(
+    values_map: dict[str, Value], node: fx.Node, loc: Location
+) -> Value:
+    """aten.replication_pad{1,2,3}d.default -> coreai.pad<replicate>."""
+    return _replace_pad_with_mode(values_map, node, loc, "replicate")
+
+
 def _conv_transpose(
     x: Value,
     weight: Value,
@@ -1522,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.
+    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])
@@ -3440,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,
@@ -3551,8 +3682,14 @@ def sdpa_maskless(q: Value, k: Value, v: Value) -> Value:
     "prod.default": replace_prod_default,
     "prod.dim_int": replace_prod_dim_int,
     "reciprocal.default": replace_reciprocal,
+    "reflection_pad1d.default": replace_reflection_pad,
+    "reflection_pad2d.default": replace_reflection_pad,
+    "reflection_pad3d.default": replace_reflection_pad,
     "relu.default": replace_unary_ops,
     "remainder.Tensor": replace_remainder,
+    "replication_pad1d.default": replace_replication_pad,
+    "replication_pad2d.default": replace_replication_pad,
+    "replication_pad3d.default": replace_replication_pad,
     "round.default": replace_unary_ops,
     "round.decimals": replace_round_decimals,
     "round": replace_unary_ops,