Skip to content

Fix conv transpose lowering#40

Merged
jakesabathia2 merged 2 commits into
apple:mainfrom
jakesabathia2:eng/PR-181169322
Jul 15, 2026
Merged

Fix conv transpose lowering#40
jakesabathia2 merged 2 commits into
apple:mainfrom
jakesabathia2:eng/PR-181169322

Conversation

@jakesabathia2

Copy link
Copy Markdown
Contributor

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:

  %22 = "coreai.shrink_dims"(%20, %21) : (tensor<?x?x?xf32>, tensor<1xsi32>) -> tensor<?x?xf32>
  RuntimeError: failed to persist mlasset

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):

  1. Static-shape erasure (the reported bug). The 1D path runs conv_transpose2d on an expanded [N,C,W,1] tensor
    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]))
  2. reshape with a runtime shape operand can't infer a static result type, so it produced tensor even
    when the conv output was fully static. A downstream squeeze (coreai.shrink_dims) then saw dynamic dims and the
    verifier rejected it.
  3. Incorrect output_padding (latent, surfaced while fixing externalize: SDPA submodule re-export drops the upper bound on the key-length dim with a static query + dynamic KV context #1). output_padding was emulated by pre-padding the
    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

  • Shape-preserving shrink-back (1D): replace the runtime-shape reshape with coreai.shrink_dims(result, [-1])
    (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).
  • Native output_padding: the Core AI conv_transpose2d op's native padding and output_pad args already match
    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

  • New regression test tests/ops/test_ops.py::test_conv_transpose_static_shape_squeeze — squeeze/reshape after
    conv_transpose1d/2d with static input, output_padding ∈ {0, 1}; validates numerics against torch eager and that
    save_asset() passes MLIR verification.
  • Updated tests/ops/test_ops_ir.py::TestConvolutionIR::test_conv_transpose2d FileCheck to the new (simpler,
    correct) IR — the op now emits 1x8x8x8 directly with no pad/slice.
  • Existing test_conv_transpose (18 cases, incl. output_padding and dynamic shapes, 1D/2D) continues to pass.

Verification

  • Radar repro (ConvTranspose1d(4,1,8,stride=2) + .squeeze(1), static input): converts and save_asset()
    succeeds, no dynamic dims in IR.
  • output_padding=1, padding=0 (the previously-wrong case): now numerically correct and static.
  • Dynamic-input transposed conv: unchanged (still uses the runtime-reshape fallback).
  • 65/65 conv tests pass; ruff clean. Full suite shows no regressions attributable to this change

Comment thread tests/ops/test_ops.py Outdated
@jakesabathia2
jakesabathia2 merged commit 698f11a into apple:main Jul 15, 2026
2 checks passed
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants