Reorganizing optim into directory structure?

[danielsuo]

Hello again! At the Princeton office, we work on, among other things, optimization algorithms for deep learning. We're interested in using flax and wanted to add some other well-known algorithms. Would you guys be open to reorganizing optim.py into a directory a la pytorch? Happy to submit a PR if so!

Usually, this would accompany a PR, but being new around here, wanted to understand how (if at all) you wanted to reorganize.

One possibility: All subclasses of OptimizerDef (except MultiOptimizer, which appears to have a circular dependency with OptimizerDef) live in their own files (e.g., Momentum, GradientDescent)

[jheek]

Sounds like a great proposal to me! And I agree it would be good to do this refactor before adding many more optimizers