diff --git a/mlx/primitives.cpp b/mlx/primitives.cpp index 6192e0759f..d4035247e7 100644 --- a/mlx/primitives.cpp +++ b/mlx/primitives.cpp @@ -1503,7 +1503,7 @@ std::vector Convolution::vjp( int64_t in_size = dilate_size(in.shape(1 + i), input_dilation_[i]); int64_t out_size = dilate_size(cotan.shape(1 + i), kernel_strides_[i]); - padding_hi[i] = safe_cast(in_size - out_size + padding_hi_[i], "conv"); + padding_hi[i] = safe_cast(in_size - out_size + padding_lo_[i], "conv"); } // Check for negative padding @@ -1515,13 +1515,21 @@ std::vector Convolution::vjp( has_neg_padding |= (pd < 0); } + // Negative padding is cropped from the output below + std::vector conv_padding_lo = padding_lo; + std::vector conv_padding_hi = padding_hi; + for (int i = 0; i < conv_padding_lo.size(); ++i) { + conv_padding_lo[i] = std::max(conv_padding_lo[i], 0); + conv_padding_hi[i] = std::max(conv_padding_hi[i], 0); + } + auto wt_trans = group_transpose(wt, 0, 1, -1); auto grad = conv_general( /* const array& input = */ cotan, /* const array& weight = */ wt_trans, /* std::vector stride = */ input_dilation_, - /* std::vector padding_lo = */ padding_lo, - /* std::vector padding_hi = */ padding_hi, + /* std::vector padding_lo = */ conv_padding_lo, + /* std::vector padding_hi = */ conv_padding_hi, /* std::vector kernel_dilation = */ kernel_dilation_, /* std::vector input_dilation = */ kernel_strides_, /* int groups = */ groups_, diff --git a/python/tests/test_conv.py b/python/tests/test_conv.py index afd4428a65..00f65eb02c 100644 --- a/python/tests/test_conv.py +++ b/python/tests/test_conv.py @@ -1185,6 +1185,35 @@ def loss_fn(kernel, inputs, strides, groups): ) self.assertEqual(grads.shape, k_shape) + def test_padding_input_grad(self): + def conv_loss(x, w, stride, padding): + out = mx.conv_general(x, w, stride=stride, padding=padding) + return mx.sum(out * out) + + def padded_loss(x, w, stride, padding): + lo, hi = padding + pad_width = [(0, 0)] + list(zip(lo, hi)) + [(0, 0)] + out = mx.conv_general(mx.pad(x, pad_width), w, stride=stride) + return mx.sum(out * out) + + for in_shape, w_shape, stride, padding in [ + ((1, 4, 1), (1, 3, 1), (1,), ([0], [1])), + ((1, 4, 1), (1, 3, 1), (2,), ([0], [1])), + ((1, 4, 1), (1, 3, 1), (1,), ([1], [0])), + ((2, 6, 6, 3), (4, 3, 3, 3), (2, 2), ([0, 1], [1, 2])), + ((2, 5, 5, 5, 2), (3, 2, 2, 2, 2), (2, 1, 2), ([0, 0, 1], [1, 1, 0])), + ((1, 4, 1), (1, 2, 1), (1,), ([3], [3])), + ((1, 4, 1), (1, 2, 1), (1,), ([0], [3])), + ((1, 4, 1), (1, 2, 1), (2,), ([3], [0])), + ((2, 5, 6, 3), (4, 2, 3, 3), (1, 2), ([4, 0], [1, 5])), + ]: + x = mx.random.normal(in_shape) + w = mx.random.normal(w_shape) + expected = mx.grad(padded_loss)(x, w, stride, padding) + grad = mx.grad(conv_loss)(x, w, stride, padding) + self.assertEqual(grad.shape, x.shape) + self.assertTrue(mx.allclose(grad, expected, atol=1e-4, rtol=1e-4)) + @unittest.skipIf(mx.cuda.is_available() and "CI" in os.environ, "flaky in CI") def test_conv_1d_with_2d(self): x = mx.random.uniform(shape=(2, 10, 16))