Repository: google/flax

Issue:
Make `ModelParamTraversal` more public?
`ModelParamTraversal` is currently somewhat hidden within `optim`. But it is much more generally useful, for example for implementing weight-decay (not as a loss) or weight standardization or spectral norm (I think).
So it seems like putting it in `traverse_util.py` (where I'd look for it) would make sense.


Discussion:
Sorry for a late reply @lucasb-eyer . This is a good proposal. Would you like to give it a shot and submit a PR? Perhaps we could move it out of optim, but keep backwards compatibility to re-exporting again in `optim`?
Currently super stretched on finishing multiple projects, so won't be able to get to it anytime in the next few weeks, sorry.
I'll give this one a try.
I'm removing myself from this issue again since I didn't find time to work on it, so if anyone would like to give this a try, please go ahead!

Resolving patch:
diff --git a/flax/optim/base.py b/flax/optim/base.py
--- a/flax/optim/base.py
+++ b/flax/optim/base.py
@@ -30,6 +30,9 @@
 
 from ..core import FrozenDict, unfreeze
 
+# Backwards compatibility symbol import.
+ModelParamTraversal = traverse_util.ModelParamTraversal
+
 
 @struct.dataclass
 class OptimizerState:
@@ -416,17 +419,6 @@ def restore_state(self, target, opt_state, state_dict):
     return self.optimizer_def.restore_state(target, opt_state, state_dict)
 
 
-def _get_params_dict(inputs):
-  if isinstance(inputs, base.Model):
-    return inputs.params
-  elif isinstance(inputs, (dict, FrozenDict)):
-    return unfreeze(inputs)
-  else:
-    raise ValueError(
-        'Can only traverse a flax Model instance or a nested dict, not '
-        f'{type(inputs)}')
-
-
 @dataclasses.dataclass
 class _ShapeDtype:
   shape: Any
@@ -442,23 +434,24 @@ def create(cls, value):
 
 
 class MultiOptimizer(OptimizerDef):
-  """ 
-  A MultiOptimizer is subclass of :class:`OptimizerDef` and useful for applying 
-  separate optimizer algorithms to various subsets of the model parameters. 
-  
-  The example below creates two optimizers using :class:`ModelParamTraversal`:
+  """
+  A MultiOptimizer is subclass of :class:`OptimizerDef` and useful for applying
+  separate optimizer algorithms to various subsets of the model parameters.
+
+  The example below creates two optimizers using
+  :class:`flax.traverse_util.ModelParamTraversal`:
   one to optimize ``kernel`` parameters and to optimize ``bias`` parameters.
   Note each optimizer is created with a different learning rate::
 
-    kernels = optim.ModelParamTraversal(lambda path, _: 'kernel' in path)
-    biases = optim.ModelParamTraversal(lambda path, _: 'bias' in path)
+    kernels = traverse_util.ModelParamTraversal(lambda path, _: 'kernel' in path)
+    biases = traverse_util.ModelParamTraversal(lambda path, _: 'bias' in path)
     kernel_opt = optim.Momentum(learning_rate=0.01)
     bias_opt = optim.Momentum(learning_rate=0.1)
     opt_def = MultiOptimizer((kernels, kernel_opt), (biases, bias_opt))
     optimizer = opt_def.create(model)
 
   In order to train only a subset of the parameters, you can simply use a single
-  :class:`ModelParamTraversal` instance.
+  :class:`flax.traverse_util.ModelParamTraversal` instance.
 
   If you want to update the learning rates of both optimizers online with
   different learning rate schedules, you should update the learning rates when
@@ -467,9 +460,9 @@ class MultiOptimizer(OptimizerDef):
 
     hparams = optimizer.optimizer_def.hyper_params
     new_optimizer = optimizer.apply_gradient(
-        grads, 
+        grads,
         hyper_params=[
-          hparams[0].replace(learning_rate=0.2), 
+          hparams[0].replace(learning_rate=0.2),
           hparams[1].replace(learning_rate=jnp.where(step < 1000, 0., lr)),
         ])
   """
@@ -546,63 +539,3 @@ def update_hyper_params(self, **hyper_param_overrides):
     if hyper_param_overrides:
       hps = [hp.replace(**hyper_param_overrides) for hp in hps]
     return hps
-
-
-def _sorted_items(x):
-  """Returns items of a dict ordered by keys."""
-  return sorted(x.items(), key=lambda x: x[0])
-
-
-class ModelParamTraversal(traverse_util.Traversal):
-  """Select model parameters using a name filter.
-  
-  This traversal operates on a nested dictionary of parameters and selects a
-  subset based on the `filter_fn` argument.
-
-  See :class:`MultiOptimizer` for an example of how to use 
-  :class:`ModelParamTraversal` to update subsets of the parameter tree with a
-  specific optimizer.
-
-  Backward compatibility:
-  When using the old api the parameters can be encapsulated in a 
-  :class:`flax.nn.Model` instance.
-  """
-
-  def __init__(self, filter_fn):
-    """Constructor a new ModelParamTraversal.
-
-    Args:
-      filter_fn: a function that takes a parameter's full name and its value and
-        returns whether this parameter should be selected or not. The name of a
-        parameter is determined by the module hierarchy and the parameter name
-        (for example: '/module/sub_module/parameter_name').
-    """
-    self._filter_fn = filter_fn
-
-  def iterate(self, inputs):
-    params = _get_params_dict(inputs)
-    flat_dict = traverse_util.flatten_dict(params)
-    for key, value in _sorted_items(flat_dict):
-      path = '/' + '/'.join(key)
-      if self._filter_fn(path, value):
-        yield value
-
-  def update(self, fn, inputs):
-    params = _get_params_dict(inputs)
-    flat_dict = traverse_util.flatten_dict(params, keep_empty_nodes=True)
-    new_dict = {}
-    for key, value in _sorted_items(flat_dict):
-      # empty_node is not an actual leave. It's just a stub for empty nodes
-      # in the nested dict.
-      if value is not traverse_util.empty_node:
-        path = '/' + '/'.join(key)
-        if self._filter_fn(path, value):
-          value = fn(value)
-      new_dict[key] = value
-    new_params = traverse_util.unflatten_dict(new_dict)
-    if isinstance(inputs, base.Model):
-      return inputs.replace(params=new_params)
-    elif isinstance(inputs, FrozenDict):
-      return FrozenDict(new_params)
-    else:
-      return new_params
diff --git a/flax/struct.py b/flax/struct.py
--- a/flax/struct.py
+++ b/flax/struct.py
@@ -12,21 +12,6 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-
-# Copyright 2020 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.
-
 """Utilities for defining custom classes that can be used with jax transformations.
 """
 
diff --git a/flax/traverse_util.py b/flax/traverse_util.py
--- a/flax/traverse_util.py
+++ b/flax/traverse_util.py
@@ -12,21 +12,6 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-
-# Copyright 2020 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.
-
 """A utility for traversing immutable datastructures.
 
 A Traversal can be used to iterate and update complex data structures.
@@ -60,12 +45,12 @@
 import dataclasses
 
 import jax
+import flax
 
 from . import struct
 
 
-# the empty node is a struct.dataclass to 
-# be compatible with JAX.
+# the empty node is a struct.dataclass to be compatible with JAX.
 @struct.dataclass
 class _EmptyNode:
   pass
@@ -381,3 +366,74 @@ def update(self, fn, inputs):
 
   def iterate(self, inputs):
     yield from jax.tree_leaves(inputs)
+
+
+def _get_params_dict(inputs):
+  if isinstance(inputs, flax.nn.Model):
+    return inputs.params
+  elif isinstance(inputs, (dict, flax.core.FrozenDict)):
+    return flax.core.unfreeze(inputs)
+  else:
+    raise ValueError(
+        'Can only traverse a flax Model instance or a nested dict, not '
+        f'{type(inputs)}')
+
+
+def _sorted_items(x):
+  """Returns items of a dict ordered by keys."""
+  return sorted(x.items(), key=lambda x: x[0])
+
+
+class ModelParamTraversal(Traversal):
+  """Select model parameters using a name filter.
+
+  This traversal operates on a nested dictionary of parameters and selects a
+  subset based on the `filter_fn` argument.
+
+  See :class:`flax.optim.MultiOptimizer` for an example of how to use
+  :class:`ModelParamTraversal` to update subsets of the parameter tree with a
+  specific optimizer.
+
+  Backward compatibility:
+  When using the old api the parameters can be encapsulated in a
+  :class:`flax.nn.Model` instance.
+  """
+
+  def __init__(self, filter_fn):
+    """Constructor a new ModelParamTraversal.
+
+    Args:
+      filter_fn: a function that takes a parameter's full name and its value and
+        returns whether this parameter should be selected or not. The name of a
+        parameter is determined by the module hierarchy and the parameter name
+        (for example: '/module/sub_module/parameter_name').
+    """
+    self._filter_fn = filter_fn
+
+  def iterate(self, inputs):
+    params = _get_params_dict(inputs)
+    flat_dict = flatten_dict(params)
+    for key, value in _sorted_items(flat_dict):
+      path = '/' + '/'.join(key)
+      if self._filter_fn(path, value):
+        yield value
+
+  def update(self, fn, inputs):
+    params = _get_params_dict(inputs)
+    flat_dict = flatten_dict(params, keep_empty_nodes=True)
+    new_dict = {}
+    for key, value in _sorted_items(flat_dict):
+      # empty_node is not an actual leave. It's just a stub for empty nodes
+      # in the nested dict.
+      if value is not empty_node:
+        path = '/' + '/'.join(key)
+        if self._filter_fn(path, value):
+          value = fn(value)
+      new_dict[key] = value
+    new_params = unflatten_dict(new_dict)
+    if isinstance(inputs, flax.nn.base.Model):
+      return inputs.replace(params=new_params)
+    elif isinstance(inputs, flax.core.FrozenDict):
+      return flax.core.FrozenDict(new_params)
+    else:
+      return new_params


Test patch:
diff --git a/tests/optim_test.py b/tests/optim_test.py
--- a/tests/optim_test.py
+++ b/tests/optim_test.py
@@ -113,58 +113,6 @@ def test_empty_optimizer(self):
     self.assertEqual(new_optimizer.state, expected_state)
 
 
-class ModelParamTraversalTest(absltest.TestCase):
-
-  def test_only_works_on_model_params(self):
-    traversal = optim.ModelParamTraversal(lambda *_: True)
-    with self.assertRaises(ValueError):
-      list(traversal.iterate([]))
-
-  def test_param_selection(self):
-    params = {
-        'x': {
-            'kernel': 1,
-            'bias': 2,
-            'y': {
-                'kernel': 3,
-                'bias': 4,
-            },
-            'z': {},
-        },
-    }
-    expected_params = {
-        'x': {
-            'kernel': 2,
-            'bias': 2,
-            'y': {
-                'kernel': 6,
-                'bias': 4,
-            },
-            'z': {}
-        },
-    }
-    names = []
-    def filter_fn(name, _):
-      names.append(name)  # track names passed to filter_fn for testing
-      return 'kernel' in name
-    traversal = optim.ModelParamTraversal(filter_fn)
-
-    # Model
-    model = nn.Model(None, params)
-    values = list(traversal.iterate(model))
-    configs = [
-      (nn.Model(None, params), nn.Model(None, expected_params)),
-      (params, expected_params),
-      (FrozenDict(params), FrozenDict(expected_params)),
-    ]
-    for model, expected_model in configs:
-      self.assertEqual(values, [1, 3])
-      self.assertEqual(set(names), set([
-          '/x/kernel', '/x/bias', '/x/y/kernel', '/x/y/bias']))
-      new_model = traversal.update(lambda x: x + x, model)
-      self.assertEqual(new_model, expected_model)
-
-
 class MultiOptimizerTest(absltest.TestCase):
 
   def test_multi_optimizer(self):
@@ -200,10 +148,10 @@ def test_multi_optimizer_multiple_matches(self):
     params = {'a': {'x': 0., 'y': 0.}, 'b': {'y': 0, 'z': 0.}}
     opt_a = optim.GradientDescent(learning_rate=1.)
     opt_b = optim.GradientDescent(learning_rate=10.)
-    t_a = optim.ModelParamTraversal(
+    t_a = traverse_util.ModelParamTraversal(
       lambda path, _: path.endswith('/x') or path.endswith('/y')
     )
-    t_b = optim.ModelParamTraversal(
+    t_b = traverse_util.ModelParamTraversal(
       lambda path, value: value.dtype == jnp.int32 or path.endswith('/z')
     )
     optimizer_def = optim.MultiOptimizer((t_a, opt_a), (t_b, opt_b))
diff --git a/tests/traverse_util_test.py b/tests/traverse_util_test.py
--- a/tests/traverse_util_test.py
+++ b/tests/traverse_util_test.py
@@ -16,11 +16,9 @@
 
 
 import collections
-
 from absl.testing import absltest
-
+import flax
 from flax import traverse_util
-
 import jax
 
 # Parse absl flags test_srcdir and test_tmpdir.
@@ -187,5 +185,58 @@ def test_flatten_dict_is_leaf(self):
     xs_restore = traverse_util.unflatten_dict(flat_xs)
     self.assertEqual(xs, xs_restore)
 
+
+class ModelParamTraversalTest(absltest.TestCase):
+
+  def test_only_works_on_model_params(self):
+    traversal = traverse_util.ModelParamTraversal(lambda *_: True)
+    with self.assertRaises(ValueError):
+      list(traversal.iterate([]))
+
+  def test_param_selection(self):
+    params = {
+        'x': {
+            'kernel': 1,
+            'bias': 2,
+            'y': {
+                'kernel': 3,
+                'bias': 4,
+            },
+            'z': {},
+        },
+    }
+    expected_params = {
+        'x': {
+            'kernel': 2,
+            'bias': 2,
+            'y': {
+                'kernel': 6,
+                'bias': 4,
+            },
+            'z': {}
+        },
+    }
+    names = []
+    def filter_fn(name, _):
+      names.append(name)  # track names passed to filter_fn for testing
+      return 'kernel' in name
+    traversal = traverse_util.ModelParamTraversal(filter_fn)
+
+    # Model
+    model = flax.nn.Model(None, params)
+    values = list(traversal.iterate(model))
+    configs = [
+      (flax.nn.Model(None, params), flax.nn.Model(None, expected_params)),
+      (params, expected_params),
+      (flax.core.FrozenDict(params), flax.core.FrozenDict(expected_params)),
+    ]
+    for model, expected_model in configs:
+      self.assertEqual(values, [1, 3])
+      self.assertEqual(set(names), set([
+          '/x/kernel', '/x/bias', '/x/y/kernel', '/x/y/bias']))
+      new_model = traversal.update(lambda x: x + x, model)
+      self.assertEqual(new_model, expected_model)
+
+
 if __name__ == '__main__':
   absltest.main()