Repository navigation
feat(ops): implement convTranspose2d — CPU kernel, backward and StableHLO lowering - #1464
Merged
Merged
Conversation
…eHLO lowering (#1371) - Promote TensorOps.convTranspose2d from a @NotImplemented default-throwing stub to an abstract @Diff op (PyTorch conv_transpose2d semantics, weight layout [C_in, C_out/groups, kH, kW]). - DefaultCpuOps: scatter-add kernel with full parameter validation (stride, padding, outputPadding < max(stride, dilation), dilation, groups, bias shape, dtype guard) — mirrors convTranspose1d, extended to NCHW. - DefaultExecutionTape: convTranspose2dBackward adjoint (input, weight and bias grads) wired into the backward dispatch table. - VoidTensorOps: shape-only implementation incl. groups-aware C_out. - StableHLO: lower convTranspose2d to stablehlo.convolution with lhs_dilate = stride, spatially flipped kernel (stablehlo.reverse) read with [i, o, 0, 1] dim numbers, and padding rewritten to low = dilation*(k-1) - padding, high = low + outputPadding. Grouped lowering is rejected explicitly (needs a kernel reshape). - TransposedConv2d: remove the @NotImplemented marker — forward now runs. - Tests: CPU kernel vs double-precision reference across the parameter grid (stride/padding/outputPadding/dilation/groups/bias + hand-computed cases), finite-difference backward parity, StableHLO lowering assertions. - Regenerate api dumps (convTranspose2d abstract, convTranspose2dBackward). Closes #1373, closes #1374, closes #1375, closes #1376. Part of #1371.
|
📖 Documentation Preview The documentation has been built successfully for this PR. Generated Files:
Artifacts:
This comment will be updated automatically when the PR is updated. |
aharakal
approved these changes
Oct 31, 2026
…transpose2d # Conflicts: # docs/modules/ROOT/pages/reference/operators/generated/index.adoc # docs/modules/ROOT/pages/reference/ops-status-matrix.adoc
|
📖 Documentation Preview The documentation has been built successfully for this PR. Generated Files:
Artifacts:
This comment will be updated automatically when the PR is updated. |
This was referenced Oct 31, 2026
Merged
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.
Summary
Implements 2D transposed convolution end to end, closing the ❌ rows for
TensorOps.convTranspose2dandTransposedConv2din the operator coverage matrix.TensorOps):convTranspose2dpromoted from a@NotImplementeddefault-throwing stub to an abstract@Diffop. PyTorchconv_transpose2dsemantics: weight[C_in, C_out/groups, kH, kW], output size(in-1)*stride - 2*padding + dilation*(k-1) + outputPadding + 1.DefaultCpuOps): scatter-add overout[b, oc, ih*sH-pH+kh*dH, iw*sW-pW+kw*dW] += in * w, mirroringconvTranspose1d, with full validation (shapes, groups divisibility,outputPadding < max(stride, dilation)per axis, bias shape, dtype guard so non-float dtypes fail loudly instead of corrupting).DefaultExecutionTape):convTranspose2dBackwardadjoint for input, weight and bias, wired into the dispatch table (the autodiff-coverage guard enforces this once the op is@Diff).NeuralNetOperationsConverter): the standard conv-transpose rewrite —stablehlo.convolutionwithlhs_dilate = stride, window stride 1, kernel spatially flipped viastablehlo.reverse dims = [2, 3]and read with[i, o, 0, 1]dim numbers, padding rewritten tolow = dilation*(k-1) - padding,high = low + outputPadding. Grouped lowering is explicitly rejected (needs a kernel reshape; tracked in [Lane 2 · kotlin-core] TransposedConv2d — StableHLO lowering for convTranspose2d #1375 follow-up if needed).VoidTensorOpsshape-only impl (groups-awareC_out),RecordingExecutiondelegation, KSP tracing-wrapper bias-nullability list, test fake inComputeGraphExecutorTest.convTranspose2dnow ✅ for cpu, Backward and StableHLO;TransposedConv2dlayer ✅.Tests
ConvTranspose2dOpsTest(backend-cpu, commonTest): hand-computed basics (k2/s1 full scatter, k2/s2 non-overlapping stamp) plus a double-precision reference across the parameter grid — multi-channel stride-2 pad-1 outputPadding-1, asymmetric stride/padding, dilation 2, groups 2, bias, andTransposedConv2d.forwardlayer integration incl.outputSize()agreement; validation failure cases.OpsAutodiffBackwardTest: finite-difference parity for the new backward, basic and strided/padded variants.ConvTranspose2dConverterTest(compile-hlo): asserts the lowering structure (reverse,lhs_dilate, pad rewrite, dim numbers) and that grouped lowering is not silently wrong.apiDumpregenerated andapiCheckgreen; all affected module suites pass locally.Closes #1373, closes #1374, closes #1375, closes #1376. Part of #1371 (lanes 1/2a/2b/2c/4 + docs partial; review lane #1378 is this PR's review).