(self)
| 13 | |
| 14 | class TestBase(mlx_tests.MLXTestCase): |
| 15 | def test_module_utilities(self): |
| 16 | m = nn.Sequential( |
| 17 | nn.Sequential(nn.Linear(2, 10), nn.relu), |
| 18 | nn.Sequential(nn.Linear(10, 10), nn.ReLU()), |
| 19 | nn.Linear(10, 1), |
| 20 | mx.sigmoid, |
| 21 | ) |
| 22 | |
| 23 | children = m.children() |
| 24 | self.assertTrue(isinstance(children, dict)) |
| 25 | self.assertEqual(len(children), 1) |
| 26 | self.assertTrue(isinstance(children["layers"], list)) |
| 27 | self.assertEqual(len(children["layers"]), 4) |
| 28 | self.assertEqual(children["layers"][3], {}) |
| 29 | flat_children = tree_flatten(children, is_leaf=nn.Module.is_module) |
| 30 | self.assertEqual(len(flat_children), 3) |
| 31 | |
| 32 | leaves = tree_flatten(m.leaf_modules(), is_leaf=nn.Module.is_module) |
| 33 | self.assertEqual(len(leaves), 4) |
| 34 | self.assertEqual(leaves[0][0], "layers.0.layers.0") |
| 35 | self.assertEqual(leaves[1][0], "layers.1.layers.0") |
| 36 | self.assertEqual(leaves[2][0], "layers.1.layers.1") |
| 37 | self.assertEqual(leaves[3][0], "layers.2") |
| 38 | self.assertTrue(leaves[0][1] is m.layers[0].layers[0]) |
| 39 | self.assertTrue(leaves[1][1] is m.layers[1].layers[0]) |
| 40 | self.assertTrue(leaves[2][1] is m.layers[1].layers[1]) |
| 41 | self.assertTrue(leaves[3][1] is m.layers[2]) |
| 42 | |
| 43 | m.eval() |
| 44 | |
| 45 | def assert_not_training(k, m): |
| 46 | self.assertFalse(m.training) |
| 47 | |
| 48 | m.apply_to_modules(assert_not_training) |
| 49 | |
| 50 | m.train() |
| 51 | |
| 52 | def assert_training(k, m): |
| 53 | self.assertTrue(m.training) |
| 54 | |
| 55 | m.apply_to_modules(assert_training) |
| 56 | |
| 57 | def test_module_attributes(self): |
| 58 | class Model(nn.Module): |
nothing calls this directly
no test coverage detected