Repository: google/flax

Issue:
Re-used BatchNorm layer with named axis can't be initialised in train mode
### Problem you have encountered:

When trying to initialise a model with a re-used `BatchNorm` layer a failure occurs when `use_running_average=False` and I've set a named axis (e.g. `axis_name="batch"`). Here is a minimal example which will fail:

```
class TestNet(nn.Module):
    @nn.compact
    def __call__(self, x, train: bool = True):
        norm = nn.BatchNorm(
            name="norm",
            use_running_average=not train,
            momentum=0.9,
            epsilon=1e-5,
            axis_name="batch"
        )

        for _ in range(2):
            x = norm(x) 

        return x

key = random.PRNGKey(0)
model = TestNet()
variables = model.init(key, jnp.ones((10,)))
```

### What you expected to happen:

I'd expect the initialization to be successful since this works if any of these three conditions are not met:
1. `use_running_average=False`,
2. there is a named axis, and
3. the `BatchNorm` layer is not reused. 

Instead, I get the following error...

### Logs, error messages, etc:

```
NameError: unbound axis name: batch. The following axis names (e.g. defined by pmap) are available to collective operations: []
```

### Steps to reproduce:
Here is a Colab to reproduce the failure as well as successful cases when any of the conditions above are not met: https://colab.research.google.com/drive/1N7Wk6eUdW4UO6Ckj_tUscDhOlO0DtY0P?usp=sharing&forceEdit=true&sandboxMode=true

Please let me know if I can clarify anything! 



Discussion:
This is indeed a known bug and it has existed for a long time I'm working on a fix now.

The reason why it hasn't been fixed before is that re-using a BatchNorm layer is rarely the correct behaviour because two inputs share batch statistics even if they aren't i.i.d..


Resolving patch:
diff --git a/flax/linen/normalization.py b/flax/linen/normalization.py
--- a/flax/linen/normalization.py
+++ b/flax/linen/normalization.py
@@ -76,6 +76,13 @@ class BatchNorm(Module):
   def __call__(self, x, use_running_average: Optional[bool] = None):
     """Normalizes the input using batch statistics.
 
+    NOTE:
+    During initialization (when parameters are mutable) the running average
+    of the batch statistics will not be updated. Therefore, the inputs
+    fed during initialization don't need to match that of the actual input
+    distribution and the reduction axis (set with `axis_name`) does not have
+    to exist.
+
     Args:
       x: the input to be normalized.
       use_running_average: if true, the statistics stored in batch_stats
@@ -93,8 +100,8 @@ def __call__(self, x, use_running_average: Optional[bool] = None):
     reduced_feature_shape = tuple(d for i, d in enumerate(x.shape) if i in axis)
     reduction_axis = tuple(i for i in range(x.ndim) if i not in axis)
 
-    # we detect if we're in initialization via empty variable tree.
-    initializing = not self.has_variable('batch_stats', 'mean')
+    # see NOTE above on initialization behavior
+    initializing = self.is_mutable_collection('params')
 
     ra_mean = self.variable('batch_stats', 'mean',
                             lambda s: jnp.zeros(s, jnp.float32),


Test patch:
diff --git a/tests/linen/linen_test.py b/tests/linen/linen_test.py
--- a/tests/linen/linen_test.py
+++ b/tests/linen/linen_test.py
@@ -146,6 +146,23 @@ def test_group_norm_raises(self):
     with self.assertRaises(ValueError):
       model_cls.init_with_output(key2, x)
 
+  def test_batch_norm_multi_init(self):
+    class Foo(nn.Module):
+      @nn.compact
+      def __call__(self, x):
+        norm = nn.BatchNorm(
+            name="norm",
+            use_running_average=False,
+            axis_name="batch",
+        )
+        x = norm(x)
+        return x, norm(x)
+
+    key = random.PRNGKey(0)
+    model = Foo()
+    x = random.normal(random.PRNGKey(1), (2, 4))
+    (y1, y2), variables = model.init_with_output(key, x)
+    np.testing.assert_allclose(y1, y2, rtol=1e-4)
 
 class StochasticTest(absltest.TestCase):