Future of flax.optim

[PhilipVinc]

Hi flax developers,

First of all I'd like to thank you for the very nice project you are developing.

Over at netket we are in the final phase of our project of bringing our quantum-physics and machine-learning framework over to JAX+flax, and we are writing some documentation and material.

As of now we are coupled to your current flax.optim interface, but before committing I would like to know what are your plans for the future.
I was reading the FLIP#1009, and it seems to me you plan to deprecate it and rely only on optax. Did I understand correctly?

Also, do you plan on updating flax.optim in order to be api-compatible with optax (namely changing OptimizerDef.create -> .init and a bunch over minor changes) or you plan on keeping different apis?

[andsteing]

Hi PhilipVinc

Yes we plan indeed to replace all our optimizers and use Optax exclusively.

The discussion for this is on PR #1011 - see especially the Update Plan section. Feel free to add any concerns you might have with the update plan directly in the FLIP so we have a central document for all the discussions.

There were some discussions about keeping a simple interface for both Optax & our previous flax.optim but we decided to keep it simple and update all existing code to use Optax directly.

[PhilipVinc]

Thanks!

So if we support the current optax api this should be enough?
Flax and Optax already work together quite well therefore most of your discussion is only about writing documentation and examples, while there won't be any API changes, other than deprecating flax.optim?

[andsteing]

Only supporting the Optax API should be future proof.

[PhilipVinc]

Out of curiosity: by looking at optax it seems that it is not very well documented (few examples, no docs, few docstrings...), very young project (version 0.0.2): optax does not look like a very mature project, even if it's definitely more powerful than flax.optim.

I therefore wonder if you also have plans to work on optax itself, or if you know that optax will receive significant contributions in the future?

[mtthss]

Hello! An optax dev here! :)

I definitely agree that there is space for improving examples and documentation in optax,
We are planning to invest time to improve both.

Regarding examples,
the planned port of Flax examples and tutorials to using optax will provide a first big improvement,
and we will also revamp the examples/ folder in optax itself to include examples of using optax with Flax,

Regarding documentation,
we are already planning to work on this,
but feel free to reach out with suggestions on anything that you think should be better documented,
the more feedback the better :)

In general, we encourage everyone in the Flax community to feel empowered to open issues on Optax, highlight your priorities, and contribute to the library! I think that consolidating on a single optimisation library shared between both current optax and current flax.optim users will really be a win win for everyone, because we will have a larger user and dev base then either project could achieve on its own

[ingmarschuster]

I just hit upon this thread by coincidence. Maybe it would make sense to highlight the plan for switching to OPTAX and deprecate flax.optim in the documentation of flax.optim flax.readthedocs.io?

[mtthss]

I think it's a good idea!
@avital

[andsteing]

This is being done as part of #1357 - feel free to comment on PR #1358