Fix MultiOptimizer crash when an optimizer's group is empty

python/mlx/optimizers/optimizers.py

@@ -193,7 +193,8 @@ def _split_dictionary(self, gradients: dict):
                     parts[i].append((k, g))
                     break
 
-        return [tree_unflatten(p) for p in parts]
+        # tree_unflatten([]) returns a list so use a dict for empty parts
+        return [tree_unflatten(p) if p else {} for p in parts]
 
     def init(self, parameters: dict):
         for o, p in zip(self.optimizers, self._split_dictionary(parameters)):

python/tests/test_optimizers.py

@@ -657,6 +657,33 @@ def test_multi_optimizer_with_parameterless_layers(self):
         self.assertFalse(mx.array_equal(w, model.layers[0].weight))
         self.assertFalse(mx.array_equal(b, model.layers[0].bias))
 
+    def test_multi_optimizer_with_empty_group(self):
+        mx.random.seed(0)
+        # no biases so all the weights are 2D
+        model = nn.Sequential(
+            nn.Linear(4, 4, bias=False), nn.ReLU(), nn.Linear(4, 4, bias=False)
+        )
+        mx.eval(model.parameters())
+
+        loss_and_grad = nn.value_and_grad(model, lambda m, x: m(x).sum())
+        _, grads = loss_and_grad(model, mx.ones((1, 4)))
+        expected = opt.Adam(learning_rate=0.01).apply_gradients(grads, model)
+
+        # SGD gets no weights, first as the fallback and then as the first optimizer
+        optimizers = [
+            opt.MultiOptimizer(
+                [opt.Adam(learning_rate=0.01), opt.SGD(learning_rate=0.01)],
+                [lambda _, w: w.ndim >= 2],
+            ),
+            opt.MultiOptimizer(
+                [opt.SGD(learning_rate=0.01), opt.Adam(learning_rate=0.01)],
+                [lambda _, w: w.ndim < 2],
+            ),
+        ]
+        for optimizer in optimizers:
+            new_params = optimizer.apply_gradients(grads, model)
+            self.assertTrue(tree_equal(mx.array_equal, expected, new_params))
+
 
 if __name__ == "__main__":
     mlx_tests.MLXTestRunner()