Repository: google/jax
Issue: BatchNorm breaks for non-default axis values

The current BatchNorm implementation breaks if the `axis` argument takes values where the missing axis is not the last one, e.g. `axis=(0, 2, 3)`, because the axes are dropped for `gamma` and `beta` and not restored later.
I guess this can be fixed by removing this line: https://github.com/google/jax/blob/7e93bff2b95e6f49be7e2cd310e8953bbbdcdc9d/jax/experimental/stax.py#L134
(and adding `tuple` in the line before)

Changed files:
jax/experimental/stax.py, tests/stax_test.py

diff --git a/jax/experimental/stax.py b/jax/experimental/stax.py
index 7b4be22a0..7382e7313 100644
--- a/jax/experimental/stax.py
+++ b/jax/experimental/stax.py
@@ -130,12 +130,16 @@ def BatchNorm(axis=(0, 1, 2), epsilon=1e-5, center=True, scale=True,
   _gamma_init = lambda shape: gamma_init(shape) if scale else ()
   axis = (axis,) if np.isscalar(axis) else axis
   def init_fun(input_shape):
-    shape = (1 if i in axis else d for i, d in enumerate(input_shape))
-    shape = tuple(itertools.dropwhile(lambda x: x == 1, shape))
+    shape = tuple(d for i, d in enumerate(input_shape) if i not in axis)
     beta, gamma = _beta_init(shape), _gamma_init(shape)
     return input_shape, (beta, gamma)
   def apply_fun(params, x, **kwargs):
     beta, gamma = params
+    # 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]
     mean, var = np.mean(x, axis, keepdims=True), fastvar(x, axis, keepdims=True)
     z = (x - mean) / np.sqrt(var + epsilon)
     if center and scale: return gamma * z + beta
diff --git a/tests/stax_test.py b/tests/stax_test.py
index 284a566fe..094c646b1 100644
--- a/tests/stax_test.py
+++ b/tests/stax_test.py
@@ -163,7 +163,7 @@ class StaxTest(jtu.JaxTestCase):
     init_fun, apply_fun = stax.FanInConcat(axis)
     _CheckShapeAgreement(self, init_fun, apply_fun, input_shapes)
 
-  def testIsuse182(self):
+  def testIssue182(self):
     init_fun, apply_fun = stax.Softmax
     input_shape = (10, 3)
     inputs = onp.arange(30.).astype("float32").reshape(input_shape)
@@ -175,5 +175,34 @@ class StaxTest(jtu.JaxTestCase):
     assert onp.allclose(onp.sum(onp.asarray(out), -1), 1.)
 
 
+  def testBatchNormShapeNHWC(self):
+    init_fun, apply_fun = stax.BatchNorm(axis=(0, 1, 2))
+    input_shape = (4, 5, 6, 7)
+    inputs = random_inputs(onp.random.RandomState(0), input_shape)
+
+    out_shape, params = init_fun(input_shape)
+    out = apply_fun(params, inputs)
+
+    self.assertEqual(out_shape, input_shape)
+    beta, gamma = params
+    self.assertEqual(beta.shape, (7,))
+    self.assertEqual(gamma.shape, (7,))
+    self.assertEqual(out_shape, out.shape)
+
+  def testBatchNormShapeNCHW(self):
+    # Regression test for https://github.com/google/jax/issues/461
+    init_fun, apply_fun = stax.BatchNorm(axis=(0, 2, 3))
+    input_shape = (4, 5, 6, 7)
+    inputs = random_inputs(onp.random.RandomState(0), input_shape)
+
+    out_shape, params = init_fun(input_shape)
+    out = apply_fun(params, inputs)
+
+    self.assertEqual(out_shape, input_shape)
+    beta, gamma = params
+    self.assertEqual(beta.shape, (5,))
+    self.assertEqual(gamma.shape, (5,))
+    self.assertEqual(out_shape, out.shape)
+
 if __name__ == "__main__":
   absltest.main()