[linen] Use stack instead of concatenate in compute_stats, to h…

…andle scalar stats case.

Added test for scalar stats, which were broken by previous change.

PiperOrigin-RevId: 549263368

flax/linen/normalization.py

@@ -95,12 +95,7 @@ def pmean(x):
 
   if use_mean:
     if use_fast_variance:
-      mean = x.mean(axes)
-      mean2 = _abs_sq(x).mean(axes)
-      if mean.ndim > 0:
-        mean, mean2 = jnp.split(pmean(jnp.concatenate([mean, mean2])), 2)
-      else:
-        mean, mean2 = pmean(mean), pmean(mean2)
+      mean, mean2 = pmean(jnp.stack([x.mean(axes), _abs_sq(x).mean(axes)]))
       # mean2 - _abs_sq(mean) is not guaranteed to be non-negative due
       # to floating point round-off errors.
       var = jnp.maximum(0.0, mean2 - _abs_sq(mean))

tests/linen/linen_test.py

@@ -178,6 +178,7 @@ def test_batch_norm_complex(self):
       {'reduction_axes': -1},
       {'reduction_axes': 1},
       {'reduction_axes': (1, 2)},
+      {'reduction_axes': (0, 1, 2)},
       {'reduction_axes': -1, 'use_fast_variance': False},
   )
   def test_layer_norm(self, reduction_axes, use_fast_variance=True):