Skip to content

[Relax] Add conv3d_transpose and ONNX ConvTranspose 3D support#18948

Merged
tlopex merged 2 commits intoapache:mainfrom
Dayuxiaoshui:main
Mar 29, 2026
Merged

[Relax] Add conv3d_transpose and ONNX ConvTranspose 3D support#18948
tlopex merged 2 commits intoapache:mainfrom
Dayuxiaoshui:main

Conversation

@Dayuxiaoshui
Copy link
Copy Markdown
Contributor

@Dayuxiaoshui Dayuxiaoshui commented Mar 28, 2026

Introduce relax.nn.conv3d_transpose (attrs, C++ inference/layout, Python API) and lower it to TOPI group_conv3d_transpose_ncdhw when using NCDHW/IODHW with dilation 1, matching the conv2d_transpose legalization policy.

Wire the Relax ONNX frontend to emit conv3d_transpose for 5D inputs. Extend tests for ONNX, struct info, LegalizeOps, and TVMScript round-trip; fix ConvTranspose test output spatial size to include output_padding.#18945

Introduce relax.nn.conv3d_transpose (attrs, C++ inference/layout, Python API) and lower it to TOPI group_conv3d_transpose_ncdhw when using NCDHW/IODHW with dilation 1, matching the conv2d_transpose legalization policy.

Wire the Relax ONNX frontend to emit conv3d_transpose for 5D inputs. Extend tests for ONNX, struct info, LegalizeOps, and TVMScript round-trip; fix ConvTranspose test output spatial size to include output_padding.
Copy link
Copy Markdown
Contributor

@gemini-code-assist gemini-code-assist bot left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Code Review

This pull request implements the conv3d_transpose operator in TVM Relax, including its C++ and Python APIs, attribute definitions, and ONNX frontend support. It also provides legalization logic to TOPI and comprehensive unit tests for shape inference and frontend integration. The review feedback identifies a typo in the Python docstring, suggests a more idiomatic dilation check in the legalization pass, and points out a redundant, unreachable logic block in the C++ struct info inference code.

- Fix conv3d_transpose out_dtype docstring wording
- Use any(d != 1 for d in dilation) in conv2d/conv3d transpose legalize
- Remove unreachable output_padding else-if in conv2d/conv3d transpose infer
@Dayuxiaoshui
Copy link
Copy Markdown
Contributor Author

cc @tlopex

Copy link
Copy Markdown
Member

@tlopex tlopex left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

LGTM

@tlopex tlopex merged commit 4de1f11 into apache:main Mar 29, 2026
9 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