BatchNorm doesn't work with center or scale turned off
BatchNorm raises TypeError when using without center or scale or both. Seems pretty straightforward to fix (e.x. as a workaround I've set beta to None `beta = None if not center else beta[ed]` if center is False, the same for gamma but there might be some better solution).

Exception:
beta = beta[ed] # ed=(None, slice(None, None, None))
TypeError: tuple indices must be integers or slices, not tuple

Code:
```
from jax import np
from jax import random
from jax.experimental import stax


init_fn, apply_fn = stax.BatchNorm(axis=0, center=False, scale=False)
_, params = init_fn(random.PRNGKey(0), (-1, 1))

x = np.zeros((2, 1))
apply_fn(params, x)
```

Details:
jax==0.1.62
jaxlib==0.1.42
diff --git a/jax/experimental/stax.py b/jax/experimental/stax.py
index d2413181c..b5881a0f6 100644
--- a/jax/experimental/stax.py
+++ b/jax/experimental/stax.py
@@ -134,12 +134,10 @@ def BatchNorm(axis=(0, 1, 2), epsilon=1e-5, center=True, scale=True,
     # TODO(phawkins): np.expand_dims should accept an axis tuple.
     # (https://github.com/numpy/numpy/issues/12290)
     ed = tuple(None if i in axis else slice(None) for i in range(np.ndim(x)))
-    beta = beta[ed]
-    gamma = gamma[ed]
     z = normalize(x, axis, epsilon=epsilon)
-    if center and scale: return gamma * z + beta
-    if center: return z + beta
-    if scale: return gamma * z
+    if center and scale: return gamma[ed] * z + beta[ed]
+    if center: return z + beta[ed]
+    if scale: return gamma[ed] * z
     return z
   return init_fun, apply_fun
 
diff --git a/tests/stax_test.py b/tests/stax_test.py
index 0c009a593..27d9dc858 100644
--- a/tests/stax_test.py
+++ b/tests/stax_test.py
@@ -212,6 +212,19 @@ class StaxTest(jtu.JaxTestCase):
     assert out_shape == out.shape
     assert onp.allclose(onp.sum(onp.asarray(out), -1), 1.)
 
+  def testBatchNormNoScaleOrCenter(self):
+    key = random.PRNGKey(0)
+    axes = (0, 1, 2)
+    init_fun, apply_fun = stax.BatchNorm(axis=axes, center=False, scale=False)
+    input_shape = (4, 5, 6, 7)
+    inputs = random_inputs(onp.random.RandomState(0), input_shape)
+
+    out_shape, params = init_fun(key, input_shape)
+    out = apply_fun(params, inputs)
+    means = onp.mean(out, axis=(0, 1, 2))
+    std_devs = onp.std(out, axis=(0, 1, 2))
+    assert onp.allclose(means, onp.zeros_like(means), atol=1e-4)
+    assert onp.allclose(std_devs, onp.ones_like(std_devs), atol=1e-4)
 
   def testBatchNormShapeNHWC(self):
     key = random.PRNGKey(0)
['jax/experimental/stax.py', 'tests/stax_test.py']