Fix sharding of quantized models with non-power-of-2 bits (#3006)

python/tests/test_nn.py

@@ -211,6 +211,18 @@ class TestBase(mlx_tests.MLXTestCase):
         size = tree_reduce(lambda acc, p: acc + p.size, qlin.trainable_parameters(), 0)
         self.assertTrue(size > 0)
 
+    def test_quantized_sharded_linear_construction(self):
+        input_dims, output_dims = 1536, 1024
+        for bits in [2, 3, 4, 5, 6, 8]:
+            lin = nn.Linear(input_dims, output_dims)
+            qlin = lin.to_quantized(bits=bits)
+
+            slin1 = nn.QuantizedAllToShardedLinear.from_quantized_linear(qlin)
+            self.assertEqual(slin1.weight.shape, qlin.weight.shape)
+
+            slin2 = nn.QuantizedShardedToAllLinear.from_quantized_linear(qlin)
+            self.assertEqual(slin2.weight.shape, qlin.weight.shape)
+
     def test_grad_of_module(self):
         class Model(nn.Module):
             def __init__(self):