Fix conv transpose lowering#40
Merged
Merged
Conversation
gokulkrishna98
approved these changes
Jul 15, 2026
jakesabathia2
force-pushed
the
eng/PR-181169322
branch
from
July 15, 2026 21:38
85a1d7e to
a4d75c7
Compare
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.
Converting a model with a squeeze/reshape after a transposed convolution (nn.ConvTranspose1d/2d) failed MLIR verification at save_asset() — even with fully static input shapes:
This blocked converting audio source-separation models (e.g. Meta's htdemucs, whose iSTFT/decoder uses transposed convolutions heavily). This PR fixes the transposed-conv lowering so static shapes are preserved end-to-end.
Root cause
Not the transposed conv op itself (its output is static). Two issues in _conv_transpose
(coreai_torch/_aten_to_core.py):
and then reshaped back to [N,C,W] using a runtime shape tensor:
coreai.reshape(result, coreai.slice_(coreai.get_shape(result), [0], [3], [1]))
when the conv output was fully static. A downstream squeeze (coreai.shrink_dims) then saw dynamic dims and the
verifier rejected it.
input, which in a stride>1 transposed conv adds stride to the output instead of output_padding (e.g. padding=0,
output_padding=1, stride=2 → output length 136 vs PyTorch's 135). This was silently wrong at runtime and only
masked at compile time by the dynamic result type.
Fix
(the true inverse of the expand_dims used to go 1D→2D) when the trailing dim is statically 1; fall back to the
runtime reshape only when the input is dynamic (where the op reports all dims as dynamic).
PyTorch semantics exactly, so the entire manual pre-pad/crop juggling was removed and padding/output_padding
are passed straight through. This is both correct and simpler (−55/+19 lines).
Tests
conv_transpose1d/2d with static input, output_padding ∈ {0, 1}; validates numerics against torch eager and that
save_asset() passes MLIR verification.
correct) IR — the op now emits 1x8x8x8 directly with no pad/slice.
Verification
succeeds, no dynamic dims in IR.