From fdde4c821c38930e2f2a19933532ed318080ca4f Mon Sep 17 00:00:00 2001 From: ishtihoss Date: Sat, 26 Sep 2026 16:19:44 -0700 Subject: [PATCH] Fix MultiOptimizer crash when an optimizer's group is empty --- python/mlx/optimizers/optimizers.py | 3 ++- python/tests/test_optimizers.py | 27 +++++++++++++++++++++++++++ 2 files changed, 29 insertions(+), 1 deletion(-) diff --git a/python/mlx/optimizers/optimizers.py b/python/mlx/optimizers/optimizers.py index 3966f961f5..cce599e11f 100644 --- a/python/mlx/optimizers/optimizers.py +++ b/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)): diff --git a/python/tests/test_optimizers.py b/python/tests/test_optimizers.py index c7caf5b954..f24a38ea36 100644 --- a/python/tests/test_optimizers.py +++ b/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()