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

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(