Updating mlx.optimizers.Adam dynamically

[chimezie]

It is not clear from the documentation, but if I wanted to update the learning rate of an mlx.optimizers.Adam instance dynamically, (to support learning rate schedules, for example), will updating the .learning_rate attribute have the desired effect?

[angeloskath]

It would have the desired effect indeed. We are looking into adding learning rate schedulers but for now one can simply change the attribute.

I recently added such an example in the transformer_lm example in mlx-examples https://github.com/ml-explore/mlx-examples/blob/7fbca214b142dbedb73d696329f154d6683e23fb/transformer_lm/main.py#L112

[chimezie]

Thanks, I'm looking forward to a framework for learning rate schedules, but I have been able to use a callback invoked every iteration that sets this attribute for now.