Skip to content

Fix conv input gradient shape with uneven or large padding - #4662

Merged
zcbenz merged 1 commit into
ml-explore:mainfrom
MarcosAsh:fix/conv-vjp-padding
Oct 10, 2026
Merged

zcbenz merged 1 commit into
ml-explore:mainfrom
MarcosAsh:fix/conv-vjp-padding

Conversation

@MarcosAsh

Copy link
Copy Markdown
Contributor
  • ☑️ I understand it is strictly prohibited to use AI to write PR description
  • AI usage disclosure:
    I used AI to help me diagnose and find this bug. All work coding and testing was done by me.

The input gradient of conv_general comes back with the wrong shape when the padding is uneven or larger than kernel - 1. In most cases you get the wrong shape sometimes it raises a confusing [conv] Spatial dimensions of input after padding cannot be smaller than weight... error from inside the backward pass. Mostly it's silent, but not always. I ran into this bug when trying out the keras-mlx backend.

x = mx.random.normal((1, 6, 2)); w = mx.random.normal((3, 3, 2))
f = lambda x: mx.conv_general(x, w, padding=([0], [1])).sum()
mx.grad(f)(x).shape  # (1, 7, 2), should be (1, 6, 2)

With the same input and kernel 3, padding (1, 1) is fine, but (2, 0) and (3, 3) both give (1, 4, 2), and (4, 0) gives (1, 0, 2).

There were two bugs in Convolution::vjp. The input-grad padding_hi was computed from padding_hi_ when it should use padding_lo_. The two are equal for symmetric padding, so it never showed up. And when the padding is larger than kernel - 1, the backward padding goes negative and was passed straight into conv_general. My proposed fix is to use padding_lo_ and clamp the backward padding to 0, letting the existing negative-padding crop handle the rest.

The new test_padding_input_grad fails on main and passes with the fix. It compares the grad against an explicit mx.pad + unpadded conv for 9 cases (uneven, large, strided, 1-D and 2-D and 3-D).

Built and tested on Linux CPU, CUDA and m4 chip (CPU).

@zcbenz
zcbenz force-pushed the fix/conv-vjp-padding branch from 6d15473 to 95c98c6 Compare October 9, 2026 10:17

@zcbenz zcbenz left a comment

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.

Thanks for the nice fix!

@zcbenz
zcbenz merged commit 0ab821b into ml-explore:main Oct 10, 2026
30 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