From d2702a4fc1c30fd334a3bc2316202d2dca3359d3 Mon Sep 17 00:00:00 2001 From: Michelle DiMarco Date: Mon, 9 Mar 2026 22:03:50 -0700 Subject: [PATCH] Fix non-strict module update with extra weights (#3214) Co-authored-by: Michelle DiMarco --- python/mlx/nn/layers/base.py | 7 +++++++ python/tests/test_nn.py | 17 +++++++++++++++++ 2 files changed, 24 insertions(+) diff --git a/python/mlx/nn/layers/base.py b/python/mlx/nn/layers/base.py index fca65d78..f72a090a 100644 --- a/python/mlx/nn/layers/base.py +++ b/python/mlx/nn/layers/base.py @@ -340,6 +340,13 @@ class Module(dict): raise ValueError(f'Module does not have parameter named "{k}".') elif isinstance(parameters, list): for i in range(len(parameters)): + if i >= len(dst): + if strict: + raise ValueError( + f"List index {i} is out of bounds for " + f"destination of length {len(dst)}." + ) + continue current_value = dst[i] new_value = parameters[i] if isinstance(current_value, mx.array): diff --git a/python/tests/test_nn.py b/python/tests/test_nn.py index ebdfe580..174823f1 100644 --- a/python/tests/test_nn.py +++ b/python/tests/test_nn.py @@ -170,6 +170,23 @@ class TestBase(mlx_tests.MLXTestCase): # Empty weights is ok if strict is false m.load_weights([], strict=False) + # Extra weights for non-existent layers are filtered when strict + # is false. Flat keys like "extra.weight" are silently dropped by + # Module.update, but nested indexed keys like "layers.1.weight" + # cause an IndexError in tree_unflatten/update without filtering. + m = nn.Sequential(nn.Linear(2, 2)) + m.load_weights( + [ + ("layers.0.weight", mx.ones((2, 2))), + ("layers.0.bias", mx.ones((2,))), + ("layers.1.weight", mx.ones((2, 2))), + ("layers.1.bias", mx.ones((2,))), + ], + strict=False, + ) + self.assertTrue(mx.array_equal(m.layers[0].weight, mx.ones((2, 2)))) + self.assertEqual(len(m.layers), 1) + def test_module_state(self): m = nn.Linear(10, 1) m.state["hello"] = "world"