From 2eacd4fa68c42c1f092ffcacac0cd0ac89930919 Mon Sep 17 00:00:00 2001 From: Teddy Tennant Date: Fri, 28 Aug 2026 12:44:43 -0400 Subject: [PATCH] fix Conv2d same padding with per axis dilation (#17805) --- test/backend/test_nn.py | 2 +- tinygrad/nn/__init__.py | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/test/backend/test_nn.py b/test/backend/test_nn.py index 75195d88e0..c12ec91266 100644 --- a/test/backend/test_nn.py +++ b/test/backend/test_nn.py @@ -135,7 +135,7 @@ class TestNN(unittest.TestCase): def test_conv2d_same_padding_large_kernel(self): self._test_conv(Conv2d, torch.nn.Conv2d, BS=16, C1=16, DIMS=[28, 33], C2=32, K=9, S=1, P='same') def test_conv2d_same_padding_with_dilation(self): - self._test_conv(Conv2d, torch.nn.Conv2d, BS=16, C1=3, DIMS=[28, 28], C2=32, K=3, S=1, P='same', D=3) + self._test_conv(Conv2d, torch.nn.Conv2d, BS=16, C1=3, DIMS=[28, 31], C2=32, K=(3,5), S=1, P='same', D=(2,3)) def test_conv2d_same_padding_invalid_stride(self): self.assertRaises(ValueError, Conv2d, in_channels=16, out_channels=32, kernel_size=2, stride=2, padding='same') diff --git a/tinygrad/nn/__init__.py b/tinygrad/nn/__init__.py index 9b663935a4..a74a4c295f 100644 --- a/tinygrad/nn/__init__.py +++ b/tinygrad/nn/__init__.py @@ -99,7 +99,7 @@ class Conv2d: if isinstance(padding, str): if padding.lower() != 'same': raise ValueError(f"Invalid padding string {padding!r}, only 'same' is supported") if stride != 1: raise ValueError("padding='same' is not supported for strided convolutions") - pad = [(d*(k-1)//2, d*(k-1) - d*(k-1)//2) for d,k in zip(make_tuple(dilation, len(self.kernel_size)), self.kernel_size[::-1])] + pad = [(d*(k-1)//2, d*(k-1) - d*(k-1)//2) for d,k in zip(make_tuple(dilation, len(self.kernel_size))[::-1], self.kernel_size[::-1])] padding = tuple(flatten(pad)) self.stride, self.dilation, self.groups, self.padding = stride, dilation, groups, padding scale = 1 / math.sqrt(in_channels * prod(self.kernel_size))