Repository: google-deepmind/optax

Issue:
Change masked wrapper to use mask_fn instead of mask
Currently, `optax.masked` accepts a boolean mask that has the same structure as (or is a prefix of) the parameters pytree. This breaks the pattern of only requiring the parameters during `init` and not before. This issue is to discuss the possibility of changing the `mask` argument to `mask_fn`, which would be a function that takes a parameter pytree as input and returns a mask that has the same/prefix structure as the params. 

One clear advantage of the `mask_fn` approach is users can define an optimizer independent of the model (as is the case with every other transformation in optax).

The current use pattern is still possible by passing in `mask=lambda _: mask` for a premade `mask` if desired.

This was originally proposed by @jheek in the [New Optimizer API](https://github.com/google/flax/pull/1011) for Flax discussion.

cc @mtthss @andsteing


Discussion:
Sounds good,

Lets make `mask` be either a nested structure or a `Callable` returning a suitable nested structure
Great! One slight issue is the `inject_hyperparams` wrapper (#48) has no way of telling whether a callable is a schedule or a fixed hyperparameter, which would make it not compose properly with this proposed `masked`.

An immediate solution that comes to mind is to add an `ignore` argument to `inject_hyperparams`, which accepts a list of hyperparameters to ignore for scheduling:

```python
inject_hyperparams(custom_optimizer, ignore=['mask'])(
    learning_rate=schedule(...), mask=mask_fn)
```

I don't think adding an optional argument is too bad for this edge case, but perhaps a more elegant/robust solution would be better. I'll think about this more before creating a PR.

Resolving patch:
diff --git a/optax/_src/wrappers.py b/optax/_src/wrappers.py
--- a/optax/_src/wrappers.py
+++ b/optax/_src/wrappers.py
@@ -29,6 +29,7 @@
 from optax._src import transform
 
 Array = jnp.ndarray
+PyTree = Any
 
 
 def flatten(
@@ -439,12 +440,21 @@ class MaskedState(NamedTuple):
 
 
 def masked(inner: transform.GradientTransformation,
-           mask: Any) -> transform.GradientTransformation:
+           mask: Union[PyTree, Callable[[transform.Params], PyTree]]
+) -> transform.GradientTransformation:
   """Mask updates so only a subset of them are computed.
 
   For example, it is common to skip weight decay for BatchNorm scale and all
   bias parameters. In many networks, these are the only parameters with only
-  one dimension. So, you may mask these out as follows:
+  one dimension. So, you may create a mask function to mask these out as
+  follows:
+
+  ```
+  mask_fn = lambda p: jax.tree_util.tree_map(lambda x: x.ndim != 1, p)
+  custom_weight_decay = optax.masked(optax.add_decayed_weights(0.001), mask_fn)
+  ```
+
+  You may alternatively create the mask pytree upfront:
 
   ```
   mask = jax.tree_util.tree_map(lambda x: x.ndim != 1, params)
@@ -456,22 +466,24 @@ def masked(inner: transform.GradientTransformation,
 
   Args:
     inner: Inner transformation to mask.
-    mask: A PyTree with the same structure as the parameters or is a prefix of
-      the parameter PyTree. The leaves should be booleans which are `True` for
-      leaves/subtrees you want to apply the transformation to, and `False` for
-      those you want to skip.
+    mask: A PyTree with the same structure as (or is a prefix of) the
+    parameters PyTree, or a Callable that returns such a pytree given the
+    parameters/updates. The leaves should be booleans which are `True` for
+    leaves/subtrees you want to apply the transformation to, and `False` for
+    those you want to skip.
 
   Returns:
     New GradientTransformation wrapping `inner`.
   """
-  flat_mask, treedef = tree_flatten(mask)
-
   def init_fn(params):
+    flat_mask, treedef = tree_flatten(mask(params) if callable(mask) else mask)
     flat_params = treedef.flatten_up_to(params)
     masked_params = [p for p, m in zip(flat_params, flat_mask) if m]
     return MaskedState(inner_state=inner.init(masked_params))
 
   def update_fn(updates, state, params=None):
+    flat_mask, treedef = tree_flatten(mask(updates) if callable(mask) else mask)
+
     # Flatten then filter out updates/params not in the mask:
     flat_updates = treedef.flatten_up_to(updates)
     masked_updates = [g for g, m in zip(flat_updates, flat_mask) if m]


Test patch:
diff --git a/optax/_src/wrappers_test.py b/optax/_src/wrappers_test.py
--- a/optax/_src/wrappers_test.py
+++ b/optax/_src/wrappers_test.py
@@ -335,25 +335,29 @@ class MaskedTest(chex.TestCase):
 
   @chex.all_variants
   @parameterized.named_parameters(
-      ('scale', lambda: transform.scale(-1.)),  # stateless test
-      ('sgd', _build_sgd),  # stateful test
+      ('scale', lambda: transform.scale(-1.), True),  # stateless test
+      ('sgd', _build_sgd, False),  # stateful test
+      ('scale', lambda: transform.scale(-1.), True),  # stateless test
+      ('sgd', _build_sgd, False),  # stateful test
   )
-  def test_masked(self, opt_builder):
+  def test_masked(self, opt_builder, use_fn):
     mask = {'a': True,
             'b': [False, True],
             'c': {'d': True, 'e': (False, True)}}
-    params = {'a': 1, 'b': [2, 3], 'c': {'d': 4, 'e': (5, 6)}}
-    input_updates = jax.tree_util.tree_map(lambda x: x/10., params)
+    mask_arg = lambda _: mask if use_fn else mask
+    params = {'a': 1., 'b': [2., 3.], 'c': {'d': 4., 'e': (5., 6.)}}
+    params = jax.tree_map(jnp.asarray, params)
+    input_updates = jax.tree_map(lambda x: x/10., params)
 
     # Negate the updates wherever the mask is True
     def masked_negate(updates):
-      return jax.tree_util.tree_multimap(
+      return jax.tree_multimap(
           lambda upd, m: -upd if m else upd, updates, mask)
     correct_updates = masked_negate(input_updates)
 
-    init_fn, update_fn = wrappers.masked(opt_builder(), mask)
+    init_fn, update_fn = wrappers.masked(opt_builder(), mask_arg)
     update_fn = self.variant(update_fn)
-    state = init_fn(params)
+    state = self.variant(init_fn)(params)
     updates, state = update_fn(input_updates, state, params)
     chex.assert_tree_all_close(updates, correct_updates)
 
@@ -370,23 +374,24 @@ def masked_negate(updates):
   def test_prefix_mask(self, opt_builder):
     """Test when the mask is a prefix of the updates PyTree."""
     mask = {'a': True, 'b': False, 'c': {'d': False, 'e': True}}
-    params = {'a': 1, 'b': {'f': 2}, 'c': {'d': 3, 'e': ([4, 5], 6)}}
-    input_updates = jax.tree_util.tree_map(lambda x: x/10., params)
+    params = {'a': 1., 'b': {'f': 2.}, 'c': {'d': 3., 'e': ([4., 5.], 6.)}}
+    params = jax.tree_map(jnp.asarray, params)
+    input_updates = jax.tree_map(lambda x: x/10., params)
 
     # Negate the updates wherever the mask (or mask parent) is True
     def _masked_sgd_on_updates(m, upd):
-      return jax.tree_util.tree_map(lambda x: -x, upd) if m else upd
-    correct_updates = jax.tree_util.tree_multimap(
+      return jax.tree_map(lambda x: -x, upd) if m else upd
+    correct_updates = jax.tree_multimap(
         _masked_sgd_on_updates, mask, input_updates)
 
     init_fn, update_fn = wrappers.masked(opt_builder(), mask)
     update_fn = self.variant(update_fn)
-    state = init_fn(params)
+    state = self.variant(init_fn)(params)
     updates, state = update_fn(input_updates, state, params)
     chex.assert_tree_all_close(updates, correct_updates)
 
     # Check repeated application, this time with no params.
-    correct_updates = jax.tree_util.tree_multimap(
+    correct_updates = jax.tree_multimap(
         _masked_sgd_on_updates, mask, correct_updates)
     updates, state = update_fn(updates, state)
     chex.assert_tree_all_close(updates, correct_updates)
@@ -397,10 +402,11 @@ def test_update_requires_params(self):
     mask = {'a': True,
             'b': [False, True],
             'c': {'d': True, 'e': (False, True)}}
-    params = {'a': 1, 'b': [2, 3], 'c': {'d': 4, 'e': (5, 6)}}
-    input_updates = jax.tree_util.tree_map(lambda x: x/10., params)
+    params = {'a': 1., 'b': [2., 3.], 'c': {'d': 4., 'e': (5., 6.)}}
+    params = jax.tree_map(jnp.asarray, params)
+    input_updates = jax.tree_map(lambda x: x/10., params)
 
-    correct_updates = jax.tree_util.tree_multimap(
+    correct_updates = jax.tree_multimap(
         lambda m, u, p: u + weight_decay * p if m else u,
         mask, input_updates, params)
 
@@ -408,14 +414,14 @@ def test_update_requires_params(self):
         transform.additive_weight_decay(weight_decay), mask)
     update_fn = self.variant(update_fn)
 
-    state = init_fn(params)
+    state = self.variant(init_fn)(params)
     updates, state = update_fn(input_updates, state, params)
     chex.assert_tree_all_close(updates, correct_updates)
 
     params = update.apply_updates(params, updates)
 
     # Test repeated application
-    new_correct_updates = jax.tree_util.tree_multimap(
+    new_correct_updates = jax.tree_multimap(
         lambda m, u, p: u + weight_decay * p if m else u,
         mask, correct_updates, params)
     updates, state = update_fn(correct_updates, state, params)
@@ -426,22 +432,41 @@ def test_empty(self, container):
     init_fn, update_fn = wrappers.masked(_build_sgd(), container())
     update_fn(container(), init_fn(container()))
 
-  @parameterized.parameters(True, False)
-  def test_tree_mismatch_fails(self, extra_key_in_mask):
+  @parameterized.parameters(
+    (False, False), (False, True), (True, False), (True, True))
+  def test_tree_mismatch_fails(self, extra_key_in_mask, use_fn):
     mask = {'a': True,
             'b': [False, True],
             'c': {'d': True, 'e': (False, True)}}
-    params = {'a': 1, 'b': [2, 3], 'c': {'d': 4, 'e': (5, 6)}}
+    mask_arg = lambda _: mask if use_fn else mask
+    params = {'a': 1., 'b': [2., 3.], 'c': {'d': 4., 'e': (5., 6.)}}
+    params = jax.tree_map(jnp.asarray, params)
 
     if extra_key_in_mask:
       mask['c']['extra'] = True
     else:
       params['c']['extra'] = 7
 
-    init_fn = wrappers.masked(_build_sgd(), mask)[0]
+    init_fn = wrappers.masked(_build_sgd(), mask_arg)[0]
     with self.assertRaises(ValueError):
       init_fn(params)
 
+  @chex.all_variants
+  def test_mask_fn(self):
+    params = {'a': jnp.ones((1, 2)), 'b': (jnp.ones((1,)), np.ones((1, 2, 3)))}
+    mask_fn = lambda p: jax.tree_map(lambda x: x.ndim > 1, p)
+    init_fn, update_fn = wrappers.masked(transform.add_decayed_weights(0.1),
+                                         mask_fn)
+    update_fn = self.variant(update_fn)
+
+    state = self.variant(init_fn)(params)
+    grads = jax.tree_map(lambda x: x*2, params)
+    updates, state = update_fn(grads, state, params)
+    np.testing.assert_allclose(updates['a'], grads['a'] + 0.1*params['a'])
+    np.testing.assert_allclose(updates['b'][0], grads['b'][0])
+    np.testing.assert_allclose(updates['b'][1],
+                               grads['b'][1] + 0.1*params['b'][1])
+
 
 class MaybeUpdateTest(chex.TestCase):
   """Tests for the maybe_update wrapper."""