Fixes grouped Conv1d/Conv2d init to use in_channels // groups

python/mlx/nn/layers/convolution.py

@@ -51,7 +51,7 @@ def __init__(
                 f"divisible by the number of groups ({groups})"
             )
 
-        scale = math.sqrt(1 / (in_channels * kernel_size))
+        scale = math.sqrt(1 / (in_channels // groups * kernel_size))
         self.weight = mx.random.uniform(
             low=-scale,
             high=scale,
@@ -132,7 +132,7 @@ def __init__(
             lambda x: (x, x) if isinstance(x, int) else x,
             (kernel_size, stride, padding),
         )
-        scale = math.sqrt(1 / (in_channels * kernel_size[0] * kernel_size[1]))
+        scale = math.sqrt(1 / (in_channels // groups * kernel_size[0] * kernel_size[1]))
         self.weight = mx.random.uniform(
             low=-scale,
             high=scale,

python/tests/test_nn.py

@@ -1,5 +1,6 @@
 # Copyright © 2023-2024 Apple Inc.
 
+import math
 import os
 import tempfile
 import unittest
@@ -1030,6 +1031,23 @@ def test_conv2d(self):
         y = c(x)
         self.assertEqual(y.shape, (4, 7, 7, 8))
 
+    def test_conv_grouped_init(self):
+        # weights are drawn from U(-s, s) with s = 1 / sqrt(fan_in), where
+        # fan_in only counts the input channels each group sees
+        mx.random.seed(0)
+        for groups in (1, 2, 4):
+            layers = [
+                nn.Conv1d(64, 64, kernel_size=3, groups=groups),
+                nn.Conv2d(64, 64, kernel_size=3, groups=groups),
+            ]
+            for layer in layers:
+                w = layer.weight
+                fan_in = math.prod(w.shape[1:])
+                bound = 1 / math.sqrt(fan_in)
+                w_max = mx.abs(w).max().item()
+                self.assertLessEqual(w_max, bound)
+                self.assertGreater(w_max, 0.95 * bound)
+
     def test_conv_transpose_extra_repr(self):
         self.assertIn(
             "kernel_size=(3, 5)",