(self)
| 392 | self.assertTrue(mx.allclose(impure_params["bias"], uncompiled_params["bias"])) |
| 393 | |
| 394 | def test_update_lr_compiled(self): |
| 395 | params = {"w": mx.ones((5, 5))} |
| 396 | grads = tree_map(lambda x: mx.ones_like(x), params) |
| 397 | optim = opt.SGD(-1.0) |
| 398 | |
| 399 | @partial(mx.compile, inputs=optim.state) |
| 400 | def update(grads): |
| 401 | return optim.apply_gradients(grads, params) |
| 402 | |
| 403 | result = update(grads) |
| 404 | self.assertTrue(mx.allclose(result["w"], mx.full((5, 5), 2.0))) |
| 405 | optim.learning_rate = -2.0 |
| 406 | result = update(grads) |
| 407 | self.assertTrue(mx.allclose(result["w"], mx.full((5, 5), 3.0))) |
| 408 | |
| 409 | |
| 410 | class TestSchedulers(mlx_tests.MLXTestCase): |
nothing calls this directly
no test coverage detected