Repository: google/flax

Issue:
Replace flax.optim with Optax in examples
See https://github.com/google/flax/blob/master/docs/flip/1009-optimizer-api.md#update-plan

The following examples need to be updated

- [x] imagenet #1251 
- [x] mnist #1302 
- [x] nlp_seq #1916
- [x] pixelcnn #1413
- [x] ppo #1404
- [x] seq2seq #1396
- [x] vae #1361 
- [x] wmt #1476
- [x] lm1b #1479


Discussion:
Run finished:
https://tensorboard.dev/experiment/w4PcKXloQMG7gXJhjskVjw/

- final test accuracy 0.7640
- total run time 5h4m

Compared to previous run using `flax.optim.Adam` (from `examples/imagenet/README.md`):
https://tensorboard.dev/experiment/iJzNKovmS0q6k5t6k5wvOw/#scalars&_smoothingWeight=0&regexInput=v100_x8_mixed_precision

- final test accuracy 0.7647
- total run time 4h51m

Note that the final test accuracy of the imagenet example fluctuates between 0.7625 and 0.7650, so the result is compatible.
Nice! Where is the branch of the code that you ran that used Optax?
Sure, it's https://github.com/andsteing/flax/tree/andsteing/issue1053
Another finished run with the Optax code:
https://tensorboard.dev/experiment/xOpycRYnT7m3inYbEuNMxw/

- final test accuracy 0.766
- total run time 5h1m

Resolving patch:
diff --git a/examples/sst2/configs/default.py b/examples/sst2/configs/default.py
--- a/examples/sst2/configs/default.py
+++ b/examples/sst2/configs/default.py
@@ -42,6 +42,5 @@ def get_config():
   config.num_epochs = 10
 
   config.seed = 0
-  config.deterministic = False
 
   return config
diff --git a/examples/sst2/train.py b/examples/sst2/train.py
--- a/examples/sst2/train.py
+++ b/examples/sst2/train.py
@@ -13,23 +13,31 @@
 # limitations under the License.
 
 """Trains an SST2 text classifier."""
-import copy
-import functools
-from typing import Any, Callable, Dict, Iterable, Mapping, Optional, Sequence, Tuple, Union
+from typing import Any, Callable, Dict, Iterable, Optional, Sequence, Tuple, Union
 
 from absl import logging
-from flax import optim
+from flax import struct
 from flax.metrics import tensorboard
+from flax.training import train_state
 import input_pipeline
 import jax
 import jax.numpy as jnp
 import ml_collections
 import models
 import numpy as np
+import optax
 import tensorflow as tf
 
 Array = jnp.ndarray
 Example = Dict[str, Array]
+TrainState = train_state.TrainState
+
+
+class Metrics(struct.PyTreeNode):
+  """Computed metrics."""
+  loss: float
+  accuracy: float
+  count: Optional[int] = None
 
 
 @jax.vmap
@@ -42,38 +50,35 @@ def sigmoid_cross_entropy_with_logits(*, labels: Array, logits: Array) -> Array:
   return relu_logits - logits * labels + jnp.log1p(jnp.exp(neg_abs_logits))
 
 
-def get_initial_params_and_state(key, model):
-  """Returns randomly initialized parameters and a fresh model state."""
+def get_initial_params(rng, model):
+  """Returns randomly initialized parameters."""
   token_ids = jnp.ones((2, 3), jnp.int32)
   lengths = jnp.ones((2,), dtype=jnp.int32)
-  variables = model.init(key, token_ids, lengths)
-  state, params = variables.pop('params')
-  return params, state
+  variables = model.init(rng, token_ids, lengths, deterministic=True)
+  return variables['params']
 
 
-def create_optimizer(params, learning_rate, beta, weight_decay):
-  """Returns a momentum optimizer."""
-  optimizer_def = optim.Momentum(
-      learning_rate=learning_rate,
-      beta=beta,
-      weight_decay=weight_decay)
-  optimizer = optimizer_def.create(params)
-  return optimizer
+def create_train_state(rng, config: ml_collections.ConfigDict, model):
+  """Create initial training state."""
+  params = get_initial_params(rng, model)
+  tx = optax.chain(
+      optax.sgd(learning_rate=config.learning_rate, momentum=config.momentum),
+      optax.additive_weight_decay(weight_decay=config.weight_decay))
+  state = TrainState.create(apply_fn=model.apply, params=params, tx=tx)
+  return state
 
 
-def compute_metrics(*, labels: Array, logits: Array) -> Dict[str, Array]:
+def compute_metrics(*, labels: Array, logits: Array) -> Metrics:
   """Computes the metrics, summed across the batch if a batch is provided."""
   if labels.ndim == 1:  # Prevent the labels from broadcasting over the logits.
     labels = jnp.expand_dims(labels, axis=1)
   loss = sigmoid_cross_entropy_with_logits(labels=labels, logits=logits)
   binary_predictions = (logits >= 0.)
   binary_accuracy = jnp.equal(binary_predictions, labels)
-  metrics = {
-      'loss': jnp.sum(loss),
-      'accuracy': jnp.sum(binary_accuracy),
-      'count': logits.shape[0]
-  }
-  return metrics
+  return Metrics(
+      loss=jnp.sum(loss),
+      accuracy=jnp.sum(binary_accuracy),
+      count=logits.shape[0])
 
 
 def model_from_config(config: ml_collections.ConfigDict):
@@ -85,110 +90,109 @@ def model_from_config(config: ml_collections.ConfigDict):
       output_size=config.output_size,
       dropout_rate=config.dropout_rate,
       word_dropout_rate=config.word_dropout_rate,
-      unk_idx=config.unk_idx,
-      deterministic=config.deterministic)
+      unk_idx=config.unk_idx)
   return model
 
 
 def train_step(
-    config: Any,
-    optimizer: optim.Optimizer,
-    model_state: Mapping[str, Any],
+    state: TrainState,
     batch: Dict[str, Array],
     rngs: Dict[str, Any],
-) -> Tuple[optim.Optimizer, Dict[str, Any], Dict[str, Any]]:
+) -> Tuple[TrainState, Metrics]:
   """Train for a single step."""
   # Make sure to get a new RNG at every step.
-  model = model_from_config(config)
-  step = optimizer.state.step
+  step = state.step
   rngs = {name: jax.random.fold_in(rng, step) for name, rng in rngs.items()}
 
   def loss_fn(params):
-    variables = {'params': params, **model_state}
-    logits, new_model_state = model.apply(
+    variables = {'params': params}
+    logits = state.apply_fn(
         variables, batch['token_ids'], batch['length'],
-        rngs=rngs, mutable=list(model_state.keys()))
+        deterministic=False,
+        rngs=rngs)
 
     labels = batch['label']
     if labels.ndim == 1:
       labels = jnp.expand_dims(labels, 1)
     loss = jnp.mean(
         sigmoid_cross_entropy_with_logits(labels=labels, logits=logits))
-    return loss, (logits, new_model_state)
+    return loss, logits
 
   grad_fn = jax.value_and_grad(loss_fn, has_aux=True)
-  value, grad = grad_fn(optimizer.target)
-  (_, (logits, new_model_state)) = value
-  optimizer = optimizer.apply_gradient(grad)
+  value, grads = grad_fn(state.params)
+  (_, logits) = value
 
+  new_state = state.apply_gradients(grads=grads)
   metrics = compute_metrics(labels=batch['label'], logits=logits)
-  return optimizer, metrics, new_model_state
+  return new_state, metrics
 
 
-def eval_step(config: Any, params: Dict[str, Any],
-              model_state: Mapping[str, Any], batch: Dict[str, Array],
-              rngs: Dict[str, Any]) -> Tuple[Dict[str, Any], Dict[str, Any]]:
+def eval_step(state: TrainState, batch: Dict[str, Array],
+              rngs: Dict[str, Any]) -> Metrics:
   """Evaluate for a single step. Model should be in deterministic mode."""
-  model = model_from_config(config)
-  variables = {'params': params, **model_state}
-  logits, new_model_state = model.apply(
+  variables = {'params': state.params}
+  logits = state.apply_fn(
       variables, batch['token_ids'], batch['length'],
-      rngs=rngs,
-      mutable=list(model_state.keys()))
+      deterministic=True,
+      rngs=rngs)
   metrics = compute_metrics(labels=batch['label'], logits=logits)
-  return metrics, new_model_state
+  return metrics
 
 
 def normalize_batch_metrics(
-        batch_metrics: Sequence[Dict[str, Any]]) -> Dict[str, Any]:
+        batch_metrics: Sequence[Metrics]) -> Metrics:
   """Consolidates and normalizes a list of per-batch metrics dicts."""
   # Here we sum the metrics that were already summed per batch.
-  metric_names = batch_metrics[0].keys()
-  summed_metrics = {
-      k: np.sum([metrics[k] for metrics in batch_metrics]) for k in metric_names
-  }
+  total_loss = np.sum([metrics.loss for metrics in batch_metrics])
+  total_accuracy = np.sum([metrics.accuracy for metrics in batch_metrics])
+  total = np.sum([metrics.count for metrics in batch_metrics])
   # Divide each metric by the total number of items in the data set.
-  total = np.float(summed_metrics.pop('count'))
-  metrics = jax.tree_map(lambda x: x.item() / total, summed_metrics)
-  return metrics
+  return Metrics(
+      loss=total_loss.item() / total, accuracy=total_accuracy.item() / total)
+
+
+def batch_to_numpy(batch: Dict[str, tf.Tensor]) -> Dict[str, Array]:
+  """Converts a batch with TF tensors to a batch of NumPy arrays."""
+  # _numpy() reuses memory, does not make a copy.
+  # pylint: disable=protected-access
+  return jax.tree_map(lambda x: x._numpy(), batch)
 
 
 def evaluate_model(
-        eval_step_fn: Callable[..., Tuple[Dict[str, Any], Dict[str, Any]]],
-        params: Dict[str, Any],
-        model_state: Mapping[str, Any],
+        eval_step_fn: Callable[..., Any],
+        state: TrainState,
         batches: Union[Iterable[Example], tf.data.Dataset],
         epoch: int,
         rngs: Optional[Dict[str, Any]] = None
-) -> Tuple[Dict[str, Any], Mapping[str, Any]]:
+) -> Metrics:
   """Evaluate a model on a dataset."""
   batch_metrics = []
   for i, batch in enumerate(batches):
-    batch = jax.tree_map(lambda x: x._numpy(), batch)  # pylint: disable=protected-access
+    batch = batch_to_numpy(batch)
     if rngs is not None:  # New RNG for each step.
       rngs = {name: jax.random.fold_in(rng, i) for name, rng in rngs.items()}
 
-    metrics, model_state = eval_step_fn(params, model_state, batch, rngs)
+    metrics = eval_step_fn(state, batch, rngs)
     batch_metrics.append(metrics)
 
   batch_metrics = jax.device_get(batch_metrics)
   metrics = normalize_batch_metrics(batch_metrics)
   logging.info('eval  epoch %03d loss %.4f accuracy %.2f', epoch,
-               metrics['loss'], metrics['accuracy'] * 100)
-  return metrics, model_state
+               metrics.loss, metrics.accuracy * 100)
+  return metrics
 
 
-def train_epoch(train_step_fn: Callable[..., Tuple[optim.Optimizer,
-                                                   Dict[str, Any], Any]],
-                optimizer: optim.Optimizer,
-                model_state: Mapping[str, Any], train_batches: tf.data.Dataset,
-                epoch: int, rngs: Optional[Dict[str, Any]] = None):
+def train_epoch(train_step_fn: Callable[..., Tuple[TrainState, Metrics]],
+                state: TrainState,
+                train_batches: tf.data.Dataset,
+                epoch: int,
+                rngs: Optional[Dict[str, Any]] = None
+                ) -> Tuple[TrainState, Metrics]:
   """Train for a single epoch."""
   batch_metrics = []
   for batch in train_batches:
-    batch = jax.tree_map(lambda x: x._numpy(), batch)  # pylint: disable=protected-access
-    optimizer, metrics, model_state = train_step_fn(
-        optimizer, model_state, batch, rngs)
+    batch = batch_to_numpy(batch)
+    state, metrics = train_step_fn(state, batch, rngs)
     batch_metrics.append(metrics)
 
   # Compute the metrics for this epoch.
@@ -196,20 +200,20 @@ def train_epoch(train_step_fn: Callable[..., Tuple[optim.Optimizer,
   metrics = normalize_batch_metrics(batch_metrics)
 
   logging.info('train epoch %03d loss %.4f accuracy %.2f', epoch,
-               metrics['loss'], metrics['accuracy'] * 100)
+               metrics.loss, metrics.accuracy * 100)
 
-  return optimizer, metrics, model_state
+  return state, metrics
 
 
 def train_and_evaluate(config: ml_collections.ConfigDict,
-                       workdir: str) -> optim.Optimizer:
+                       workdir: str) -> TrainState:
   """Execute model training and evaluation loop.
 
   Args:
     config: Hyperparameter configuration for training and evaluation.
     workdir: Directory where the tensorboard summaries are written to.
   Returns:
-    The trained optimizer.
+    The final train state that includes the trained parameters.
   """
   # Prepare datasets.
   train_dataset = input_pipeline.TextDataset(
@@ -225,28 +229,17 @@ def train_and_evaluate(config: ml_collections.ConfigDict,
       shuffle_seed=config.seed)
   eval_batches = eval_dataset.get_batches(batch_size=config.batch_size)
 
-  # Prepare configs.
+  # Keep track of vocab size in the config so that the embedder knows it.
   config.vocab_size = len(train_dataset.vocab)
-  eval_config = copy.deepcopy(config)
-  eval_config.deterministic = True
 
   # Compile step functions.
-  train_step_fn = jax.jit(functools.partial(train_step, config))
-  eval_step_fn = jax.jit(functools.partial(eval_step, eval_config))
+  train_step_fn = jax.jit(train_step)
+  eval_step_fn = jax.jit(eval_step)
 
-  # Initialize parameters.
+  # Create model and a state that contains the parameters.
   rng = jax.random.PRNGKey(config.seed)
-  init_model = model_from_config(eval_config)
-  params, model_state = get_initial_params_and_state(rng, init_model)
-  del init_model
-  
-  # Remove intermediates for training. Otherwise our model state will fill up
-  # with intermediate outputs (exported using self.sow() commands). This will
-  # cause model_state to have a new shape on each step, triggering a new trace.
-  model_state, _ = model_state.pop('intermediates')  
-  
-  optimizer = create_optimizer(
-      params, config.learning_rate, config.momentum, config.weight_decay)
+  model = model_from_config(config)
+  state = create_train_state(rng, config, model)
 
   summary_writer = tensorboard.SummaryWriter(workdir)
   summary_writer.hparams(dict(config))
@@ -258,24 +251,23 @@ def train_and_evaluate(config: ml_collections.ConfigDict,
     # Train for one epoch.
     rng, epoch_rng = jax.random.split(rng)
     rngs = {'dropout': epoch_rng}
-    optimizer, train_metrics, model_state = train_epoch(
-        train_step_fn, optimizer, model_state, train_batches, epoch, rngs)
+    state, train_metrics = train_epoch(
+        train_step_fn, state, train_batches, epoch, rngs)
 
     # Evaluate current model on the validation data.
-    eval_metrics, _ = evaluate_model(
-        eval_step_fn, optimizer.target, model_state, eval_batches, epoch)
+    eval_metrics = evaluate_model(eval_step_fn, state, eval_batches, epoch)
 
     # Write metrics to TensorBoard.
-    summary_writer.scalar('train_loss', train_metrics['loss'], epoch)
+    summary_writer.scalar('train_loss', train_metrics.loss, epoch)
     summary_writer.scalar(
         'train_accuracy',
-        train_metrics['accuracy'] * 100,
+        train_metrics.accuracy * 100,
         epoch)
-    summary_writer.scalar('eval_loss', eval_metrics['loss'], epoch)
+    summary_writer.scalar('eval_loss', eval_metrics.loss, epoch)
     summary_writer.scalar(
         'eval_accuracy',
-        eval_metrics['accuracy'] * 100,
+        eval_metrics.accuracy * 100,
         epoch)
 
   summary_writer.flush()
-  return optimizer
+  return state


Test patch:
diff --git a/examples/sst2/train_test.py b/examples/sst2/train_test.py
new file mode 100644
--- /dev/null
+++ b/examples/sst2/train_test.py
@@ -0,0 +1,56 @@
+# Copyright 2021 The Flax Authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Tests for sst2.train."""
+from absl.testing import absltest
+from absl.testing import parameterized
+from configs import default as default_config
+import jax
+import jax.test_util
+import numpy as np
+import train
+
+# Parse absl flags test_srcdir and test_tmpdir.
+jax.config.parse_flags_with_absl()
+
+
+class TrainTest(parameterized.TestCase):
+
+  def test_train_step_updates_parameters(self):
+    """Tests if the train step updates the parameters in train state."""
+    # Create model and a state that contains the parameters.
+    config = default_config.get_config()
+    config.vocab_size = 13
+    rng = jax.random.PRNGKey(config.seed)
+    model = train.model_from_config(config)
+    state = train.create_train_state(rng, config, model)
+
+    token_ids = np.array([[2, 4, 3], [2, 6, 3]], dtype=np.int32)
+    lengths = np.array([2, 3], dtype=np.int32)
+    labels = np.zeros_like(lengths)
+    batch = {'token_ids': token_ids, 'length': lengths, 'label': labels}
+    rngs = {'dropout': rng}
+    train_step_fn = jax.jit(train.train_step)
+    new_state, metrics = train_step_fn(state, batch, rngs)
+    self.assertIsInstance(new_state, train.TrainState)
+    self.assertIsInstance(metrics, train.Metrics)
+    old_param_values = jax.tree_leaves(state.params)
+    new_param_values = jax.tree_leaves(new_state.params)
+    for old_array, new_array in zip(old_param_values, new_param_values):
+      # Make sure parameters were updated.
+      self.assertFalse(np.allclose(old_array, new_array))
+
+
+if __name__ == '__main__':
+  absltest.main()