diff --git a/coreai_torch/_aten_to_core.py b/coreai_torch/_aten_to_core.py index 75e596c..11e9530 100644 --- a/coreai_torch/_aten_to_core.py +++ b/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 diff --git a/tests/ops/test_ops.py b/tests/ops/test_ops.py index a012e4a..4a15547 100644 --- a/tests/ops/test_ops.py +++ b/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(