File size: 1,607 Bytes
712dbf0 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 |
# Copyright © 2023 Apple Inc.
import unittest
import mlx.core as mx
import mlx.nn as nn
import mlx.utils
import mlx_tests
class TestTreeUtils(mlx_tests.MLXTestCase):
def test_tree_map(self):
tree = {"a": 0, "b": 1, "c": 2}
tree = mlx.utils.tree_map(lambda x: x + 1, tree)
expected_tree = {"a": 1, "b": 2, "c": 3}
self.assertEqual(tree, expected_tree)
def test_tree_flatten(self):
tree = [{"a": 1, "b": 2}, "c"]
vals = (1, 2, "c")
flat_tree = mlx.utils.tree_flatten(tree)
self.assertEqual(list(zip(*flat_tree))[1], vals)
self.assertEqual(mlx.utils.tree_unflatten(flat_tree), tree)
def test_merge(self):
t1 = {"a": 0}
t2 = {"b": 1}
t = mlx.utils.tree_merge(t1, t2)
self.assertEqual({"a": 0, "b": 1}, t)
with self.assertRaises(ValueError):
mlx.utils.tree_merge(t1, t1)
with self.assertRaises(ValueError):
mlx.utils.tree_merge(t, t1)
mod1 = nn.Sequential(nn.Linear(2, 2), nn.Linear(2, 2))
mod2 = nn.Sequential(nn.Linear(2, 2), nn.Linear(2, 2))
mod = nn.Sequential(mod1, mod2)
params1 = {"layers": [mod1.parameters()]}
params2 = {"layers": [None, mod2.parameters()]}
params = mlx.utils.tree_merge(params1, params2)
for (k1, v1), (k2, v2) in zip(
mlx.utils.tree_flatten(params), mlx.utils.tree_flatten(mod.parameters())
):
self.assertEqual(k1, k2)
self.assertTrue(mx.array_equal(v1, v2))
if __name__ == "__main__":
mlx_tests.MLXTestRunner()
|