`stax.BatchNorm` incorrect normalization
jax.experimental.stax on current master, line 140 ( `BatchNorm` function):
```python
z = (x - mean) / (var + epsilon)**2
```
The square in the denominator should probably be a square root.
diff --git a/jax/experimental/stax.py b/jax/experimental/stax.py
index 311c9d3f5..caff9a984 100644
--- a/jax/experimental/stax.py
+++ b/jax/experimental/stax.py
@@ -137,7 +137,7 @@ def BatchNorm(axis=(0, 1, 2), epsilon=1e-5, center=True, scale=True,
   def apply_fun(params, x, rng=None):
     beta, gamma = params
     mean, var = np.mean(x, axis, keepdims=True), fastvar(x, axis, keepdims=True)
-    z = (x - mean) / (var + epsilon)**2
+    z = (x - mean) / np.sqrt(var + epsilon)
     if center and scale: return gamma * z + beta
     if center: return z + beta
     if scale: return gamma * z
['jax/experimental/stax.py']