_aten_to_core, _utils: harden mixed-source SymInt lowerings under dynamic shapes

Six related fixes that surface together when exporting models whose FX
graphs combine SymInt-derived shape arithmetic with mixed source
element types and ranks:

1. _aten_to_core_resolver / replace_binary_ops _op_map: register a
   bare 'pow' entry alongside the variant-suffixed ones. Some
   torch.export rewrites leave ``aten.pow`` as the OpOverloadPacket
   target with no overload suffix; without this entry the converter
   raises ``Unsupported ATen op: pow``.
2. Same registries for bare 'round': torch.export can leave
   ``aten.round`` without a ``.default`` overload, mirroring the pow
   case.
3. upsample_build_output_shape_dynamic: ensure each (out_h, out_w)
   operand is rank-1 with int32 element type before the concat that
   builds the output shape — the dialect verifier rejects mixed-rank
   / mixed-element-type concat inputs. Hits when out_h/out_w are
   SymInts derived from ``round(SymFloat)`` arithmetic.
4. get_operand mixed-list path (SymInt + plain int): normalise each
   resolved Value to the same canonical rank-1 si32 form and emit
   plain-int constants with explicit ``dtype=np.int32`` so the
   dim-vector concat sees uniform operands. Hits ops like ``view``,
   ``expand``, ``reshape``, ``repeat`` whenever a dim list mixes
   SymInts with ints.
5. replace_cat: when one input has a dynamic non-concat axis and a
   sibling has a known static size for that axis, reshape the
   dynamic side to that static size before the concat. Localised
   shape inference using the fact that all non-concat dims must be
   equal — multiple distinct static sizes is left for the dialect
   verifier to reject.
6. replace_arange_start_step: unify start/end/step element types to
   the FX node's output dtype before ``coreai.range_``. Mirrors
   aten.arange's internal type promotion since coreai.range_'s
   verifier requires uniform element types.

Adds a shared ``to_rank1_int32(v)`` helper in ``_utils.py`` so fixes
3 and 4 share one canonical normalization (rank-0 → rank-1, cast to
signed int32 if needed); both call sites collapse to one line per
operand.

One regression test per non-trivial fix, each verified to FAIL
without the fix and PASS with it (verified by reverting each fix
individually):

- TestRound: bare ``aten.round`` overload-packet target must lower.
- TestUpsampleNearest2d / TestUpsampleBilinear2d::test_round_symfloat_size:
  ``round((num / aspect) ** 0.5) * 14`` output_size produces SymInts
  whose Value type doesn't match the int32 constants used elsewhere.
  Pre-fix: ``coreai.concat`` raises ``Operation creation failed``.
- TestView::test_view_with_round_symfloat_dims: same trigger applied
  to a ``view([1, C, h, w])`` mixed list. Pre-fix: ``expected the
  same element type for all inputs to concat``.
- TestCat::test_dynamic_vs_static_non_concat_axis: ``Dim.AUTO`` on
  one side + static sibling forces non-concat-axis promotion.
- TestArange::test_symint_end_with_float_start_step: float ``arange``
  with SymInt-derived end exercises the element-type unify.

Author: gokulkrishna98

coreai_torch/_aten_to_core.py

@@ -644,17 +644,20 @@ def replace_arange_start_step(
         else coreai.constant(1, dtype=start.type.element_type)
     )
 
-    # Squeeze rank-1 inputs to 0D scalars for coreai.range_.
+    # 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.
+    target_type = get_output_element_type_from_node(node)
+
     def to_scalar(v: Value) -> Value:
         if v.type.rank > 0:
-            return coreai.shrink_dims(v, list(range(v.type.rank)))
+            v = coreai.shrink_dims(v, list(range(v.type.rank)))
+        if v.type.element_type != target_type:
+            v = coreai.cast(v, target_type)
         return v
 
-    result = coreai.range_(to_scalar(start), to_scalar(end), to_scalar(step))
-    target_type = get_output_element_type_from_node(node)
-    if result.type.element_type != target_type:
-        result = coreai.cast(result, target_type)
-    return result
+    return coreai.range_(to_scalar(start), to_scalar(end), to_scalar(step))
 
 
 def replace_batch_norm(
@@ -814,6 +817,7 @@ def replace_binary_ops(
         "pow.Scalar": coreai.broadcasting_pow,
         "pow.Tensor_Tensor": coreai.broadcasting_pow,
         "pow.Tensor_Scalar": coreai.broadcasting_pow,
+        "pow": coreai.broadcasting_pow,
         "sub.Tensor": coreai.broadcasting_sub,
         "sub.Scalar": coreai.broadcasting_sub,
         "sub": coreai.broadcasting_sub,
@@ -944,6 +948,36 @@ def replace_cat(values_map: dict[str, Value], node: fx.Node, loc: Location) -> V
 
     rank = inputs[0].type.rank
     dim = dim + rank if dim < 0 else dim
+
+    # coreai.concat requires all non-concat dims to be provably equal across
+    # inputs. Under dynamic shapes, one branch can carry a dynamic non-concat
+    # axis while a sibling has a static size for the same axis — the dynamic
+    # side must in fact equal that static size, but the type system doesn't
+    # know it. Reshape such inputs to the known static size before the concat.
+    # Multiple distinct static sizes on one axis is a real mismatch and is
+    # left for the dialect verifier to reject.
+    dyn = ShapedType.get_dynamic_size()
+
+    def known_static(axis: int) -> int | None:
+        if axis == dim:
+            return None
+        sizes = {inp.type.shape[axis] for inp in inputs if inp.type.shape[axis] != dyn}
+        return next(iter(sizes)) if len(sizes) == 1 else None
+
+    statics = [known_static(a) for a in range(rank)]
+    promoted: list[Value] = []
+    for inp in inputs:
+        new_shape = [
+            statics[a]
+            if statics[a] is not None and inp.type.shape[a] == dyn
+            else inp.type.shape[a]
+            for a in range(rank)
+        ]
+        if new_shape != list(inp.type.shape):
+            inp = coreai.reshape(inp, new_shape)
+        promoted.append(inp)
+    inputs = promoted
+
     return coreai.concat(dim, inputs)
 
 
@@ -2595,6 +2629,7 @@ def replace_unary_ops(
         "log.default": coreai.log,
         "relu.default": coreai.relu,
         "round.default": coreai.round_,
+        "round": coreai.round_,
         "rsqrt.default": coreai.rsqrt,
         "sigmoid.default": coreai.sigmoid,
         "silu.default": coreai.silu,
@@ -3434,13 +3469,15 @@ def sdpa_maskless(q: Value, k: Value, v: Value) -> Value:
     "pow.Scalar": replace_binary_ops,
     "pow.Tensor_Scalar": replace_binary_ops,
     "pow.Tensor_Tensor": replace_binary_ops,
+    "pow": replace_binary_ops,
     "prod.default": replace_prod_default,
     "prod.dim_int": replace_prod_dim_int,
     "reciprocal.default": replace_reciprocal,
     "relu.default": replace_unary_ops,
     "remainder.Tensor": replace_remainder,
     "round.default": replace_unary_ops,
     "round.decimals": replace_round_decimals,
+    "round": replace_unary_ops,
     "repeat.default": replace_repeat,
     "rsqrt.default": replace_unary_ops,
     "scaled_dot_product_attention.default": replace_sdpa,

coreai_torch/_utils.py

@@ -175,6 +175,23 @@ def __exit__(self, *exc: Any) -> None:
         self.close()
 
 
+def to_rank1_int32(v: Value) -> Value:
+    """Coerce a SymInt-derived Value to canonical rank-1 si32 form.
+
+    Dim-vector concats (used to build shape operands for ``coreai.reshape``,
+    ``coreai.interpolate``, etc.) require all inputs to share rank and
+    element type. SymInt values can arrive rank-0 (e.g. from
+    ``aten._local_scalar_dense``) or with a different int variant. This
+    helper produces the form that aligns with ``coreai.constant([i],
+    dtype=np.int32)`` and ``replace_sym_size_int``.
+    """
+    if v.type.rank == 0:
+        v = coreai.reshape(v, [1])
+    if v.type.element_type != IntegerType.get_signed(32):
+        v = coreai.cast(v, np.int32)
+    return v
+
+
 def upsample_build_output_shape_dynamic(
     x: Value, out_h: int | Value, out_w: int | Value
 ) -> Value:
@@ -192,8 +209,8 @@ def upsample_build_output_shape_dynamic(
     )
     shape = coreai.cast(coreai.get_shape(x), dtype=np.int32)
     non_spatial = coreai.slice_(shape, [0], [2], [1])
-    h = [out_h] if isinstance(out_h, int) else out_h
-    w = [out_w] if isinstance(out_w, int) else out_w
+    h = [out_h] if isinstance(out_h, int) else to_rank1_int32(out_h)
+    w = [out_w] if isinstance(out_w, int) else to_rank1_int32(out_w)
     return coreai.concat(0, [non_spatial, h, w])
 
 
@@ -986,9 +1003,13 @@ def get_operand(
     if isinstance(arg, fx.Node):
         return values_map[arg.name]
     if isinstance(arg, list) and any(isinstance(e, fx.Node) for e in arg):
-        # Mixed list: resolve fx.Node elements via values_map, keep ints as constants.
+        # Mixed list: SymInt fx.Nodes + plain ints. Concat the two sources
+        # into a single rank-1 si32 dim vector. Both branches must produce
+        # the same canonical form so the concat verifier accepts them.
         dim_vals = [
-            values_map[e.name] if isinstance(e, fx.Node) else coreai.constant([e])
+            to_rank1_int32(values_map[e.name])
+            if isinstance(e, fx.Node)
+            else coreai.constant([e], dtype=np.int32)
             for e in arg
         ]
         return coreai.concat(0, dim_vals) if len(dim_vals) > 1 else dim_vals[0]