From 2c63d3b5baae8495b93c0e47ca90b4dcec9d31c7 Mon Sep 17 00:00:00 2001 From: Pujitha Paladugu <10557236+pujitha24@users.noreply.github.com> Date: Sat, 31 Oct 2026 09:25:05 -0700 Subject: [PATCH] [FIX][Relax] Restore logical layout before flattening repeat in ConvertLayout Motivation: For relax.repeat with axis=None, InferLayoutRepeat kept the producer's layout on the input. After ConvertLayout turned a conv2d into NHWC, the repeat flattened the NHWC tensor, so the flattened element order differed from the original NCHW module (reported max abs diff 16.1 on LLVM and CUDA). Approach: When axis is not given, request the initial (logical) layout for the input so a permute_dims back to the original layout is inserted before the flattening repeat. The axis-specified path is unchanged. Impact: Only the axis=None form of repeat is affected; it now produces the same flattened order as before layout conversion, at the cost of one extra permute_dims. Validation (macOS arm64, built from source without LLVM/CUDA): - Before the fix, test_conv2d_repeat_flatten passed against the buggy IR (no permute before repeat). After the fix that expectation fails and the updated expectation (permute_dims NHWC->NCHW, then repeat) passes. - python -m pytest tests/python/relax/test_transform_convert_layout.py -> 75 passed. - pre-commit run on the changed files: passes (clang-format reflowed one line). Not run: numeric before/after execution of the issue reproducer (no LLVM or CUDA available here); correctness rests on the IR now matching the logical NCHW order before the flatten. Report: https://github.com/apache/tvm/issues/20826 Signed-off-by: Pujitha Paladugu <10557236+pujitha24@users.noreply.github.com> Assisted-by: claude-sonnet-5-5 (via Claude Code) --- src/relax/op/tensor/manipulate.cc | 6 ++++-- tests/python/relax/test_transform_convert_layout.py | 5 ++++- 2 files changed, 8 insertions(+), 3 deletions(-) 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