Skip to content

implement Convolution::output_shape#2601

Merged
awni merged 1 commit intoml-explore:mainfrom
josharian:shapeless-conv
Sep 22, 2025
Merged

implement Convolution::output_shape#2601
awni merged 1 commit intoml-explore:mainfrom
josharian:shapeless-conv

Conversation

@josharian
Copy link
Copy Markdown
Contributor

@josharian josharian commented Sep 17, 2025

Proposed changes

Implement Convolution::output_shape so that convolutions can be used with shapeless function exports.

Please review carefully: My C++ is not super strong, and I'm new to mlx.

The code in mlx/conv_shape.h is 100% code movement from mlx/ops.cpp, no changes aside from marking conv_out_shape as inline.

Updates #2599

Checklist

  • I have read the CONTRIBUTING document
  • I have run pre-commit run --all-files to format my code / installed pre-commit prior to committing changes
  • I have added tests that prove my fix is effective or that my feature works
  • I have updated the necessary documentation (if needed)

mlx/conv_shape.h Outdated
@@ -0,0 +1,121 @@
// Copyright © 2025 Apple Inc.
Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

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

We don't want a new file for this. Instead you can make it static on the Convolution primitive and then just use it in the ops.

Comment on lines +1465 to +1476
auto expand = [spatial_dims](const std::vector<int>& v, int def) {
if (v.empty()) {
return std::vector<int>(spatial_dims, def);
}
if (v.size() == 1) {
return std::vector<int>(spatial_dims, v[0]);
}
if (v.size() == spatial_dims) {
return v;
}
throw std::invalid_argument("Convolution param rank mismatch");
};
Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

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

I think these parameters should already be expanded once they are in the primitive right?

- pull conv_out_shape out for re-use
- add Conv::output_shape
- add e2e python tests confirming shapeless=True support and correctness

Updates ml-explore#2599
@josharian
Copy link
Copy Markdown
Contributor Author

Great feedback. Thanks! Fixed and re-tested and pushed a new commit. PTAL.

Copy link
Copy Markdown
Member

@awni awni left a comment

Choose a reason for hiding this comment

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

Looks great, thanks!

@awni awni merged commit aa9d44b into ml-explore:main Sep 22, 2025
7 of 8 checks passed
@josharian josharian deleted the shapeless-conv branch September 22, 2025 23:49
faisalmemon pushed a commit to faisalmemon/mlx that referenced this pull request Oct 30, 2025
- pull conv_out_shape out for re-use
- add Conv::output_shape
- add e2e python tests confirming shapeless=True support and correctness

Updates ml-explore#2599
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