Add dilation parameter and tests

Author: aretor

fft_conv_pytorch/fft_conv.py

@@ -130,6 +130,7 @@ def __init__(
         kernel_size: Union[int, Iterable[int]],
         padding: Union[int, Iterable[int]] = 0,
         stride: Union[int, Iterable[int]] = 1,
+        dilation: Union[int, Iterable[int]] = 1,
         groups: int = 1,
         bias: bool = True,
         ndim: int = 1,
@@ -150,6 +151,7 @@ def __init__(
         self.kernel_size = kernel_size
         self.padding = padding
         self.stride = stride
+        self.dilation = dilation
         self.groups = groups
         self.use_bias = bias
 
@@ -165,9 +167,22 @@ def __init__(
             )
 
         kernel_size = to_ntuple(kernel_size, ndim)
-        self.weight = nn.Parameter(
-            torch.randn(out_channels, in_channels // groups, *kernel_size)
+        dilation = to_ntuple(dilation, ndim)
+        total_size = tuple(
+            ((ks - 1) * dil + 1)
+            for ks, dil in zip(kernel_size, dilation)
         )
+        weight = torch.zeros(out_channels, in_channels // groups, *total_size)
+        fill = torch.randn(out_channels, in_channels // groups, *kernel_size)
+        ids = tuple(
+            torch.arange(0, tot_sz, dil)
+            for tot_sz, dil in zip(total_size, dilation)
+        )
+        
+        # workaround bc PyTorch doesn't support [:, :, <tensor tuple>] indexing
+        weight[(slice(None), slice(None),) + torch.meshgrid(*ids)] = fill
+        
+        self.weight = nn.Parameter(weight)
         self.bias = nn.Parameter(torch.randn(out_channels)) if bias else None
 
     def forward(self, signal):

tests/test_fft_conv.py

@@ -27,6 +27,7 @@ def _gcd(x: int, y: int) -> int:
 @pytest.mark.parametrize("kernel_size", [1, 2, 3])
 @pytest.mark.parametrize("padding", [0, 1])
 @pytest.mark.parametrize("stride", [1, 2, 3])
+@pytest.mark.parametrize("dilation", [1, 2, 3])
 @pytest.mark.parametrize("bias", [True, False])
 @pytest.mark.parametrize("ndim", [1, 2, 3])
 @pytest.mark.parametrize("input_size", [7, 8])
@@ -36,6 +37,7 @@ def test_fft_conv(
     kernel_size: Union[int, Iterable[int]],
     padding: Union[int, Iterable[int]],
     stride: Union[int, Iterable[int]],
+    dilation: Union[int, Iterable[int]],
     groups: int,
     bias: bool,
     ndim: int,
@@ -49,6 +51,7 @@ def test_fft_conv(
         kernel_size=kernel_size,
         padding=padding,
         stride=stride,
+        dilation=dilation,
         groups=groups,
         bias=bias,
         ndim=ndim,