Repository: google/jax

Issue:
Adding regularizers to Stax
Some utilities for regularizing the cost functions would be useful in Stax. I have not seen this mentioned in #137. Is this planned?

Also, I was wondering about the best way to implement this in pure JAX. For example considering L2 regularization, with the flattening utility from #190 I could write:

```
w_flat, _ = ravel_pytree(w)
reg = np.sum(w_flat ** 2.0)
```

Looking a little bit on the code of tree_util I can also write:

```
reg= tree_reduce(lambda x, y: np.sum(x) + np.sum(y), tree_map(lambda wi: np.sum(wi**2.0), w))
```

Is there a simpler solution I am missing? More in general it would be useful to document tree utilities in the main doc.


Discussion:
Thanks for suggesting this!

The `tree_reduce` version is better because otherwise all the arrays have to be concatenated together. It's possible that XLA could optimize away that explicit concatenation when under an `@jit`, but I'm not sure it will.

Here's another way to write it which avoids traversing the tree twice:

```python
from jax.tree_util import tree_flatten
import jax.numpy as np

def l2_squared(pytree):
  leaves, _ = tree_flatten(pytree)
  return sum(np.vdot(x, x) for x in leaves)
```

We could add this to Stax, or perhaps to an MNIST example file. We'd like to avoid making Stax into a monolith; part of the point is to show how easy it is to write neural net "layers" libraries, so you can write one yourself with whatever features you want.
A notebook with some "advanced" use cases would be terrific! Especially on the parts of the library that are less documented at the moment.

Resolving patch:
diff --git a/jax/experimental/optimizers.py b/jax/experimental/optimizers.py
--- a/jax/experimental/optimizers.py
+++ b/jax/experimental/optimizers.py
@@ -79,7 +79,8 @@
 import jax.numpy as np
 from jax.util import partial, safe_zip, safe_map, unzip2
 from jax import tree_util
-from jax.tree_util import tree_flatten, tree_unflatten, register_pytree_node
+from jax.tree_util import (tree_map, tree_flatten, tree_unflatten,
+                           register_pytree_node)
 
 map = safe_map
 zip = safe_zip
@@ -376,3 +377,17 @@ def make_schedule(scalar_or_schedule):
     return constant(scalar_or_schedule)
   else:
     raise TypeError(type(scalar_or_schedule))
+
+
+### utilities
+
+def l2_norm(tree):
+  """Compute the l2 norm of a pytree of arrays. Useful for weight decay."""
+  leaves, _ = tree_flatten(tree)
+  return np.sqrt(sum(np.vdot(x, x) for x in leaves))
+
+def clip_grads(grad_tree, max_norm):
+  """Clip gradients stored as a pytree of arrays to maximum norm `max_norm`."""
+  norm = l2_norm(grad_tree)
+  normalize = lambda g: np.where(norm < max_norm, g, g * (max_norm / norm))
+  return tree_map(normalize, grad_tree)


Test patch:
diff --git a/tests/optimizers_test.py b/tests/optimizers_test.py
--- a/tests/optimizers_test.py
+++ b/tests/optimizers_test.py
@@ -20,6 +20,8 @@
 import functools
 
 from absl.testing import absltest
+import numpy as onp
+
 import jax.numpy as np
 import jax.test_util as jtu
 from jax import jit, grad
@@ -213,6 +215,26 @@ def get_params(opt_state):
     opt_state = init_fun(np.zeros(3))
     self.assertRaises(TypeError, lambda: update_fun(opt_state))
 
+  def testUtilityNorm(self):
+    x0 = (np.ones(2), (np.ones(3), np.ones(4)))
+    norm = optimizers.l2_norm(x0)
+    expected = onp.sqrt(onp.sum(onp.ones(2+3+4)**2))
+    self.assertAllClose(norm, expected, check_dtypes=False)
+
+  def testUtilityClipGrads(self):
+    g = (np.ones(2), (np.ones(3), np.ones(4)))
+    norm = optimizers.l2_norm(g)
+
+    ans = optimizers.clip_grads(g, 1.1 * norm)
+    expected = g
+    self.assertAllClose(ans, expected, check_dtypes=False)
+
+    ans = optimizers.l2_norm(optimizers.clip_grads(g, 0.9 * norm))
+    expected = 0.9 * norm
+    self.assertAllClose(ans, expected, check_dtypes=False)
+
+
+
 
 if __name__ == '__main__':
   absltest.main()