Open buttercutter opened 1 year ago
As highlighted by @jcole75 , Lion is already implemented in deepmind's optax repo , so pip install git+https://github.com/deepmind/optax.git
would do the work. See also https://optax.readthedocs.io/en/latest/api.html#lion
As for Adam#Lion
, when I try to study how the optimizer function call returns, it seems that it is not that straightforward to translate the learning rate grafting to optax
jax
Any advice on this ?
I am also not sure yet how to implement lr grafting yet.
I noticed that there are no other choices of optimizers other than scale_by_adafactor()
This github issue would serve as a placeholder for other optimizers such as
Adam#Lion
or others in the future.TODO:
Adam#Lion
jax code