feat - Implemented Lion Optimizer and Tests scripts

[raghulchandramouli]

Hey as per the issue i have raised an pull request for the implementation of lion_optimizer, please check the code and Merge

Thanks and Regards,
RaghulChandramouli

[google-cla]

Thanks for your pull request! It looks like this may be your first contribution to a Google open source project. Before we can look at your pull request, you'll need to sign a Contributor License Agreement (CLA).

View this failed invocation of the CLA check for more information.

For the most up to date status, view the checks section at the bottom of the pull request.

[vroulet]

Check this https://optax.readthedocs.io/en/latest/api/optimizers.html#optax.lion first. Let me know tif you want to change the current implementation.

[vroulet]

Thanks ! Sorry for the delay. Do you think it would be possible to add the smooth option to the current implementation? That way we minimize code duplication and maintenance burden while getting a neat new feature.

Happy to review it then ! (ping me when it's ready)

[raghulchandramouli]

I have raised a new branch and added all the changes
https://github.com/google-deepmind/optax/pull/1497

Please track this @vroulet

[rdyro]

https://github.com/google-deepmind/optax/pull/1497 does not address Vincent's comments. See my comment there.