diff --git a/src/relax/op/tensor/manipulate.cc b/src/relax/op/tensor/manipulate.cc index fb3141b14721..48cf570d0c44 100644 --- a/src/relax/op/tensor/manipulate.cc +++ b/src/relax/op/tensor/manipulate.cc @@ -1980,9 +1980,11 @@ InferLayoutOutput InferLayoutRepeat( existing_layout = LayoutDecision(InitialLayout(ndim)); } - // When axis is not specified, the output is 1D (flattened) + // When axis is not specified, the output is the flattened input. Flattening is order-sensitive, + // so the input must be restored to its initial (logical) layout. if (!attrs->axis.has_value()) { - return InferLayoutOutput({existing_layout}, {InitialLayoutDecision(1)}, call->attrs); + return InferLayoutOutput({InitialLayoutDecision(ndim)}, {InitialLayoutDecision(1)}, + call->attrs); } // Transform the axis based on the layout diff --git a/tests/python/relax/test_transform_convert_layout.py b/tests/python/relax/test_transform_convert_layout.py index 5aa5cfc40a3e..e29aff5b21b2 100644 --- a/tests/python/relax/test_transform_convert_layout.py +++ b/tests/python/relax/test_transform_convert_layout.py @@ -5076,7 +5076,10 @@ def main( out_layout="NHWC", out_dtype="float32", ) - gv2: R.Tensor((5408,), dtype="float32") = R.repeat(gv, repeats=1) + lv2: R.Tensor((2, 4, 26, 26), dtype="float32") = R.permute_dims( + gv, axes=[0, 3, 1, 2] + ) + gv2: R.Tensor((5408,), dtype="float32") = R.repeat(lv2, repeats=1) R.output(gv2) return gv2