chore: change function with a destination dictonary object

docs/src/python/optimizers.rst

@@ -51,14 +51,14 @@ the saved state. Here's a simple example:
    optimizer.update(model, grads)
 
    # Save the state
-   state = tree_flatten(optimizer.state)
-   mx.save_safetensors("optimizer.safetensors", dict(state))
+   state = tree_flatten(optimizer.state, destination={})
+   mx.save_safetensors("optimizer.safetensors", state)
 
    # Later on, for example when loading from a checkpoint,
    # recreate the optimizer and load the state
    optimizer = optim.Adam(learning_rate=1e-2)
 
-   state = tree_unflatten(list(mx.load("optimizer.safetensors").items()))
+   state = tree_unflatten(mx.load("optimizer.safetensors"))
    optimizer.state = state
 
 Note, not every optimizer configuation parameter is saved in the state. For

docs/src/usage/export.rst

@@ -151,7 +151,7 @@ parameters, pass them as inputs to the ``call`` wrapper:
      model.update(tree_unflatten(list(params.items())))
      return model(x)
 
-   params = dict(tree_flatten(model.parameters()))
+   params = tree_flatten(model.parameters(), destination={})
    mx.export_function("model.mlxfn", call, (mx.zeros(4),), params)
 
 

python/mlx/nn/layers/base.py

@@ -178,7 +178,7 @@ class Module(dict):
 
         if strict:
             new_weights = dict(weights)
-            curr_weights = dict(tree_flatten(self.parameters()))
+            curr_weights = tree_flatten(self.parameters(), destination={})
             if extras := (new_weights.keys() - curr_weights.keys()):
                 num_extra = len(extras)
                 extras = ",\n".join(sorted(extras))
@@ -212,7 +212,7 @@ class Module(dict):
         - ``.npz`` will use :func:`mx.savez`
         - ``.safetensors`` will use :func:`mx.save_safetensors`
         """
-        params_dict = dict(tree_flatten(self.parameters()))
+        params_dict = tree_flatten(self.parameters(), destination={})
 
         if file.endswith(".npz"):
             mx.savez(file, **params_dict)

python/tests/test_nn.py

@@ -30,6 +30,7 @@ class TestBase(mlx_tests.MLXTestCase):
         self.assertEqual(len(flat_children), 3)
 
         leaves = tree_flatten(m.leaf_modules(), is_leaf=nn.Module.is_module)
+        if isinstance(leaves, list):
             self.assertEqual(len(leaves), 4)
             self.assertEqual(leaves[0][0], "layers.0.layers.0")
             self.assertEqual(leaves[1][0], "layers.1.layers.0")
@@ -80,7 +81,7 @@ class TestBase(mlx_tests.MLXTestCase):
                 self.weights = {"w1": mx.zeros((2, 2)), "w2": mx.ones((2, 2))}
 
         model = DictModule()
-        params = dict(tree_flatten(model.parameters()))
+        params = tree_flatten(model.parameters(), destination={})
         self.assertEqual(len(params), 2)
         self.assertTrue(mx.array_equal(params["weights.w1"], mx.zeros((2, 2))))
         self.assertTrue(mx.array_equal(params["weights.w2"], mx.ones((2, 2))))