Repository navigation
Conversation
…rtLayout 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: apache#20826 Signed-off-by: Pujitha Paladugu <10557236+pujitha24@users.noreply.github.com> Assisted-by: claude-sonnet-5-5 (via Claude Code)
This branch has not been deployed
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
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):
(no permute before repeat). After the fix that expectation fails and the
updated expectation (permute_dims NHWC->NCHW, then repeat) passes.
-> 75 passed.
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: #20826
Signed-off-by: Pujitha Paladugu 10557236+pujitha24@users.noreply.github.com
Assisted-by: claude-sonnet-5-5 (via Claude Code)
Fixes #20826