Removes flax.optim & bumps version to 0.6.0

This fixes #2273 PiperOrigin-RevId: 467869152
google · Aug 16, 2022 · cda7a4c · cda7a4c
1 parent 47418b3
commit cda7a4c
Show file tree

Hide file tree

Showing 21 changed files with 21 additions and 1,724 deletions.
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -6,8 +6,6 @@ vNext
 (Add your change to a random empty line to avoid merge conflicts)
 -
 -
-- Switched to using `jax.named_scope` for all profile naming, cut some pointless
-  stack traces out.
 -
 -
 -
@@ -27,6 +25,14 @@ vNext
 -
 -
 
+0.6.0
+-----
+
+- Removed deprecated optimizers in `flax.optim` package.
+- Moved `flax.optim.dynamic_scale` to `flax.training.dynamic_scale`.
+- Switched to using `jax.named_scope` for all profile naming, cut some pointless
+  stack traces out.
+
 0.5.3
 -----
 New features:

diff --git a/docs/advanced_topics/optax_update_guide.rst b/docs/advanced_topics/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:

diff --git a/docs/conf.py b/docs/conf.py
@@ -129,6 +129,7 @@
 myst_enable_extensions = ['dollarmath']
 nb_execution_excludepatterns = [
   'getting_started.ipynb', # <-- times out 
+  'optax_update_guide.ipynb', # <-- requires flax<=0.5.3
 ]
 # raise exceptions on execution so CI can catch errors
 nb_execution_allow_errors = False

diff --git a/docs/notebooks/optax_update_guide.ipynb b/docs/notebooks/optax_update_guide.ipynb
@@ -33,7 +33,8 @@
    },
    "outputs": [],
    "source": [
-    "!pip install -q flax optax"
+    "# flax.optim was deprecated after 0.5.3\n",
+    "!pip install -q --force-reinstall flax==0.5.3 optax"
    ]
   },
   {

diff --git a/docs/notebooks/optax_update_guide.md b/docs/notebooks/optax_update_guide.md
@@ -27,7 +27,8 @@ https://flax.readthedocs.io/en/latest/howtos/optax_update_guide.html
 :id: I4PiwrnnO6Fw
 :tags: [skip-execution]
 
-!pip install -q flax optax
+# flax.optim was deprecated after 0.5.3
+!pip install -q --force-reinstall flax==0.5.3 optax
 ```
 
 ```{code-cell}

diff --git a/flax/__init__.py b/flax/__init__.py
@@ -15,10 +15,12 @@
 
 """Flax API."""
 
-from . import core as core
-from . import linen as linen
+from . import core
+from . import jax_utils
+from . import linen
+from . import serialization
+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
diff --git a/flax/optim/README.md b/flax/optim/README.md
diff --git a/flax/optim/__init__.py b/flax/optim/__init__.py
diff --git a/flax/optim/adabelief.py b/flax/optim/adabelief.py
diff --git a/flax/optim/adadelta.py b/flax/optim/adadelta.py