Removes flax.optim & bumps version to 0.6.0

This fixes #2273

PiperOrigin-RevId: 459485412

CHANGELOG.md

@@ -26,6 +26,11 @@ vNext
 -
 -
 
+0.6.0
+-----
+
+- Removes deprecated optimizers in `flax.optim` package.
+
 0.5.2
 -----
 - Fixes missing PyYAML dependecy.

docs/howtos/optax_update_guide.rst

@@ -7,7 +7,7 @@ Upgrading my Codebase to Optax
 We have proposed to replace :py:mod:`flax.optim` with `Optax
 <https://optax.readthedocs.io>`_ in 2021 with `FLIP #1009
 <https://github.com/google/flax/blob/main/docs/flip/1009-optimizer-api.md>`_ and
-the Flax optimizers are now *effectively deprecated*. This guide is targeted
+the Flax optimizers have been removed in v0.6.0 - this guide is targeted
 towards :py:mod:`flax.optim` users to help them update their code to Optax.
 
 See also Optax's quick start documentation:

flax/__init__.py

@@ -15,11 +15,13 @@
 
 """Flax API."""
 
-from . import core as core
-from . import linen as linen
-from . import optim as optim
+from . import core
+from . import jax_utils
+from . import linen
+from . import serialization
+from . import struct
+from . import traverse_util
+
 # DO NOT REMOVE - Marker for internal deprecated API.
 # DO NOT REMOVE - Marker for internal logging.
 from .version import __version__
-
-# Allow `import flax`; `flax.optim.[...]`, etc