Fix conv2d_weight_grad_groups

[laggui]

Checklist

• Confirmed that run-checks all script has been executed.

Related Issues/PRs

Fixes #1877

Changes

Fixed conv2d backward weight slice assign dimensions issue when using a kernel with groups=channels and stride>1.

Testing

Added a unit test to capture this issue.

Reference values obtained from this script:

import torch
from torch import nn, Tensor

if __name__ == "__main__":
    channels = 4
    height = width = 4
    kernel_size = 3
    conv = nn.Conv2d(
        channels, channels, kernel_size, stride=2, padding=1, groups=channels
    )

    for param in conv.parameters():
        if param.ndim == 4:
            param.data = nn.Parameter(
                torch.arange(channels * kernel_size * kernel_size)
                .float()
                .reshape(channels, 1, kernel_size, kernel_size)
            )
        else:
            param.data = nn.Parameter(torch.arange(channels).float())

    numel = channels * height * width
    x = torch.arange(numel).reshape(1, channels, height, width).float().requires_grad_()
    y: Tensor = conv(x)
    grads = y.backward(torch.ones_like(y))

    torch.set_printoptions(precision=8)
    print(x.grad)
    print(conv.weight.grad)
    print(conv.bias.grad)

tensor([[[[  4.,   8.,   4.,   5.],
          [  8.,  16.,   8.,  10.],
          [  4.,   8.,   4.,   5.],
          [  7.,  14.,   7.,   8.]],

         [[ 13.,  26.,  13.,  14.],
          [ 26.,  52.,  26.,  28.],
          [ 13.,  26.,  13.,  14.],
          [ 16.,  32.,  16.,  17.]],

         [[ 22.,  44.,  22.,  23.],
          [ 44.,  88.,  44.,  46.],
          [ 22.,  44.,  22.,  23.],
          [ 25.,  50.,  25.,  26.]],

         [[ 31.,  62.,  31.,  32.],
          [ 62., 124.,  62.,  64.],
          [ 31.,  62.,  31.,  32.],
          [ 34.,  68.,  34.,  35.]]]])
tensor([[[[  5.,  10.,  12.],
          [ 10.,  20.,  24.],
          [ 18.,  36.,  40.]]],


        [[[ 21.,  42.,  44.],
          [ 42.,  84.,  88.],
          [ 50., 100., 104.]]],


        [[[ 37.,  74.,  76.],
          [ 74., 148., 152.],
          [ 82., 164., 168.]]],


        [[[ 53., 106., 108.],
          [106., 212., 216.],
          [114., 228., 232.]]]])
tensor([4., 4., 4., 4.])

[codecov]

Codecov Report

All modified and coverable lines are covered by tests ✅

Project coverage is 86.08%. Comparing base (5252440) to head (ad5e291).

Additional details and impacted files

@@            Coverage Diff             @@
##             main    #1891      +/-   ##
==========================================
+ Coverage   86.07%   86.08%   +0.01%     
==========================================
  Files         780      780              
  Lines       91416    91489      +73     
==========================================
+ Hits        78685    78758      +73     
  Misses      12731    12731              

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

[nathanielsimard]

LGTM