Repository: google-deepmind/optax

Issue:
Nesting 'masked' wrappers gives a confusing error
The _masked_ wrapper requires the model parameters to be input as a pytree. It flattens that pytree before calling init on its inner optimizer. So if you wrap masked around a optimizer that also uses masked, like lamb, optax throws a mismatched type error. It expects to receive a pytree, but instead receives a flat list of parameters. In the short term, I recommend we throw a useful error when this occurs, since it's relatively easy to mitigate this through the use of chain. In the long term I recommend rebuilding the tree in masked before passing it to the inner optimizer.


Discussion:

Resolving patch:
diff --git a/optax/_src/wrappers.py b/optax/_src/wrappers.py
--- a/optax/_src/wrappers.py
+++ b/optax/_src/wrappers.py
@@ -322,35 +322,25 @@ def masked(
   Returns:
     New GradientTransformation wrapping ``inner``.
   """
+  def mask_pytree(pytree, mask_tree):
+    return tree_map(lambda m, p: p if m else None, mask_tree, pytree)
+
   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]
+    mask_tree = mask(params) if callable(mask) else mask
+    masked_params = mask_pytree(params, mask_tree)
     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]
-
-    if params is not None:
-      flat_params = treedef.flatten_up_to(params)
-      masked_params = [p for p, m in zip(flat_params, flat_mask) if m]
-    else:
-      masked_params = None
+    mask_tree = mask(updates) if callable(mask) else mask
+    masked_updates = mask_pytree(updates, mask_tree)
+    masked_params = None if params is None else mask_pytree(params, mask_tree)
 
-    # Compute new updates
     new_masked_updates, new_inner_state = inner.update(
         masked_updates, state.inner_state, masked_params)
 
-    # Incorporate new_masked_updates into flat_updates, then unflatten
-    new_masked_updates = iter(new_masked_updates)
-    for i, m in enumerate(flat_mask):
-      if m: flat_updates[i] = next(new_masked_updates)
-
-    new_updates = treedef.unflatten(flat_updates)
+    new_updates = tree_map(
+      lambda m, new_u, old_u: new_u if m else old_u,
+      mask_tree, new_masked_updates, updates)
     return new_updates, MaskedState(inner_state=new_inner_state)
 
   return base.GradientTransformation(init_fn, update_fn)


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
@@ -14,6 +14,8 @@
 # ==============================================================================
 """Tests for `wrappers.py`."""
 
+import copy
+
 from absl.testing import absltest
 from absl.testing import parameterized
 
@@ -396,6 +398,44 @@ def test_mask_fn(self):
     np.testing.assert_allclose(updates['b'][1],
                                grads['b'][1] + 0.1*params['b'][1])
 
+  @chex.all_variants
+  @parameterized.named_parameters(
+      ('sgd', _build_sgd),
+      ('stateful_sgd', _build_stateful_sgd),
+  )
+  def test_nested_mask(self, opt_builder):
+    # https://github.com/deepmind/optax/issues/271
+    params = {'linear_1': {'w': jnp.zeros((1, 1)), 'b': jnp.zeros(1)},
+              'linear_2': {'w': jnp.zeros((1, 2)), 'b': jnp.zeros(2)},
+              'linear_3': {'w': jnp.zeros((2, 3)), 'b': jnp.zeros(3)}}
+
+    outer_mask = lambda p: jax.tree_map(lambda x: x.ndim > 1, p)
+    inner_mask = jax.tree_map(lambda _: True, params)
+    inner_mask['linear_2'] = False
+
+    inner = wrappers.masked(opt_builder(), inner_mask)
+    init_fn, update_fn = wrappers.masked(inner, outer_mask)
+
+    input_updates = jax.tree_map(jnp.ones_like, params)
+    correct_updates = copy.deepcopy(input_updates)
+    correct_updates['linear_1']['w'] *= -1.0
+    correct_updates['linear_3']['w'] *= -1.0
+
+    state = self.variant(init_fn)(params)
+    updates, state = self.variant(update_fn)(input_updates, state, params)
+    chex.assert_trees_all_close(updates, correct_updates)
+
+  @chex.all_variants
+  def test_masked_state_structure(self):
+    # https://github.com/deepmind/optax/issues/271
+    params = {'a': [jnp.ones(1), (jnp.ones(2), jnp.ones(3))],
+              'b': {'c': jnp.ones(4), 'd': jnp.ones(5)}}
+    mask = {'a': [True, (True, False)], 'b': False}
+    tx = wrappers.masked(_build_stateful_sgd(), mask)
+    trace = self.variant(tx.init)(params).inner_state[0].trace
+    expected_trace = {'a': [jnp.zeros(1), (jnp.zeros(2), None)], 'b': None}
+    chex.assert_tree_all_equal_structs(trace, expected_trace)
+
 
 class MaybeUpdateTest(chex.TestCase):
   """Tests for the maybe_update wrapper."""