[converter] Fix cat lowering when promoted shape still has dynamic axes (#14)

* _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.

* _aten_to_core: extend replace_cat dynamic-axis promotion to handle still-dynamic siblings

The previous replace_cat fix uses ``coreai.reshape(inp, new_shape)`` with
a Python list to promote a dynamic non-concat axis to its known static
size before the concat. The list form materializes the shape as an
``int32`` constant tensor, so every entry must be a real int — there is
no slot for the dynamic-size sentinel.

When a cat input has both a promotable axis (sibling has a known static
size) and an axis that is dynamic on every input (no static sibling),
post-promotion ``new_shape`` is a mix of concrete ints and the dynamic
sentinel; passing it to the list-form reshape raises
``OverflowError: Python integer -9223372036854775808 out of bounds for
int32``.

Split ``replace_cat`` into two reshape paths:

- All axes static post-promotion: keep the list-form reshape.
- Some axes still dynamic post-promotion: build the shape vector at
  runtime — ``coreai.get_shape(inp)`` for the still-dynamic axes,
  ``coreai.constant`` slices for the promoted axes, concat along axis 0
  to get a rank-1 ``int32`` Value, and pass that to Value-form
  ``coreai.ReshapeOp`` with a partially-static result type.

Adds a numerical regression that exercises the mixed path: a 4-D cat
where one input has a sibling-promotable axis and another axis that
stays dynamic on every input. Verified to fail with the list-form-only
fix and pass with this delta.

---------

Co-authored-by: gokulkrishna98 <gokulkrishna98@users.noreply.github.com>

Author: gokulkrishna98

coreai_torch/_aten_to_core.py

@@ -974,7 +974,26 @@ def known_static(axis: int) -> int | None:
             for a in range(rank)
         ]
         if new_shape != list(inp.type.shape):
-            inp = coreai.reshape(inp, new_shape)
+            if all(s != dyn for s in new_shape):
+                # All axes static post-promotion: list-form reshape packs
+                # the shape into an int32 constant tensor.
+                inp = coreai.reshape(inp, new_shape)
+            else:
+                # Mixed static / dynamic post-promotion: build the shape
+                # vector at runtime by mixing the input's actual sizes
+                # (via coreai.get_shape) for the still-dynamic axes with
+                # constants for the promoted axes.
+                runtime_shape = coreai.cast(coreai.get_shape(inp), dtype=np.int32)
+                parts = [
+                    coreai.constant([s], dtype=np.int32)
+                    if s != dyn
+                    else coreai.slice_(runtime_shape, [a], [a + 1], [1])
+                    for a, s in enumerate(new_shape)
+                ]
+                result_type = RankedTensorType.get(new_shape, inp.type.element_type)
+                inp = coreai.ReshapeOp(
+                    inp, coreai.concat(0, parts), results=[result_type]
+                ).result
         promoted.append(inp)
     inputs = promoted
 

tests/ops/test_ops.py

@@ -1024,6 +1024,32 @@ def forward(self, a: Tensor, b: Tensor) -> Tensor:
             dynamic_shapes=({1: torch.export.Dim.AUTO}, {}),
         )
 
+    async def test_partial_static_promotion_with_dynamic_axes(self) -> None:
+        """A cat input has one non-concat axis that is statically known
+        via a sibling AND another non-concat axis that is dynamic on every
+        input. After promoting the first axis to its static size, the
+        second axis remains dynamic, so the lowering must build the
+        reshape's shape vector at runtime."""
+
+        class CatModel(nn.Module):
+            def forward(self, a: Tensor, b: Tensor) -> Tensor:
+                return torch.cat([a, b], dim=2)
+
+        a = torch.rand(2, 4, 5, 6)
+        b = torch.rand(2, 4, 7, 6)
+        # Mark dim 1 of `a` dynamic (sibling `b` has static 4 there → must be
+        # promoted) and dim 0 of both inputs dynamic (no static sibling →
+        # remains dynamic post-promotion, forcing the runtime-shape path).
+        await validate_numerical_output(
+            model=CatModel().eval(),
+            a=a,
+            b=b,
+            dynamic_shapes=(
+                {0: torch.export.Dim.AUTO, 1: torch.export.Dim.AUTO},
+                {0: torch.export.Dim.AUTO},
+            ),
+        )
+
 
 @pytest.mark.parametrize("dynamic", [False, True])
 @pytest.mark.parametrize(