(self)
| 52 | self.assertEqual(optim.state, {0: 1}) |
| 53 | |
| 54 | def test_optimizers(self): |
| 55 | params = { |
| 56 | "first": [mx.zeros((10,)), mx.zeros((1,))], |
| 57 | "second": mx.zeros((1,)), |
| 58 | } |
| 59 | grads = tree_map(lambda x: mx.ones_like(x), params) |
| 60 | |
| 61 | for optim_class in optimizers_dict.values(): |
| 62 | optim = optim_class(0.1) |
| 63 | update = optim.apply_gradients(grads, params) |
| 64 | mx.eval(update) |
| 65 | equal_shape = tree_map(lambda x, y: x.shape == y.shape, params, update) |
| 66 | all_equal = all(v for _, v in mlx.utils.tree_flatten(equal_shape)) |
| 67 | self.assertTrue(all_equal) |
| 68 | |
| 69 | def test_types_conserved(self): |
| 70 | params = {"w": mx.ones((5, 5), mx.float16)} |
nothing calls this directly
no test coverage detected