Support LR schedulers (#334)

* Add a few LR schedulers * Move parents's constructor call to the top * Fix docstring * refactor optimizers into two files * add docs * nit * Fix Callable type annotation for python 3.8 --------- Co-authored-by: Awni Hannun <[email protected]> Co-authored-by: Angelos Katharopoulos <[email protected]>
2024-02-15 11:26:20 -08:00
parent 85143fecdd
commit 818cda16bc
10 changed files with 235 additions and 47 deletions
@@ -0,0 +1,20 @@
+.. _common_optimizers:
+
+Common Optimizers
+=================
+
+.. currentmodule:: mlx.optimizers
+
+.. autosummary::
+   :toctree: _autosummary
+   :template: optimizers-template.rst
+
+   SGD
+   RMSprop
+   Adagrad
+   Adafactor
+   AdaDelta
+   Adam
+   AdamW
+   Adamax
+   Lion