Repository: keras-team/keras-core

Issue:
BatchNorm layer appears broken in multi-backend Keras
classic MNIST training:
- [Keras multi-backend with BN on JAX](https://colab.research.google.com/drive/1ZZI6UK9umgCDSIRcPqVUl6YpLmRdwMF3?resourcekey=0-8YS9usjjxUuh0FUl_JW0Ow&usp=sharing). Final accuracy __19%__!!!
- [Keras multi-backend w/o BN on JAX](https://colab.research.google.com/drive/16Yj63zZHmH2oQNBtNYmP5W0KJJ_RnLQi?resourcekey=0-Ut-aMYTyo3i36iJeQB4LRQ&usp=sharing). Final accuracy: 99%
- [Keras multi-backend with BN on TF](https://colab.research.google.com/drive/1ipSnQuGe2w4KGphG6hMeWdzt3H8m36cx?resourcekey=0-ZH-RRNE3rfXqqurX8yE_YA&usp=sharing). Final accuracy: __14%__!!!
- [Keras multi-backend w/o BN on TF](https://colab.research.google.com/drive/1beSnSJDvBwqYXzaGpvdVH9bUgCM2Jtsi?resourcekey=0-jmZOshPvFOUqWOEGq-z2Dg&usp=sharing). Final accuracy: 99%
- [Keras baseline (regular tf.keras on TF, with BN)](https://colab.research.google.com/drive/1dI3RTfVvYdYER09Q6onYIdsh3vkRmEci?resourcekey=0-FgN08JmB4pWDpsL8aePs8A&usp=sharing). Final accuracy: 99%

The common factors triggering the issue seem to be: using the BatchNorm layer, on Keras multi-backend. The choice of backend does not seem to change the behavior.


Discussion:
This issue is new. A couple of weeks ago, when I made go/keras-multi-backend-demo, BatchNormalization was working normally in Keras multi-backend (i tested on JAX at the time)

Resolving patch:
diff --git a/keras_core/layers/normalization/batch_normalization.py b/keras_core/layers/normalization/batch_normalization.py
--- a/keras_core/layers/normalization/batch_normalization.py
+++ b/keras_core/layers/normalization/batch_normalization.py
@@ -201,6 +201,10 @@ def call(self, inputs, training=None, mask=None):
             self.moving_mean.assign(
                 self.moving_mean * self.momentum + mean * (1.0 - self.momentum)
             )
+            self.moving_variance.assign(
+                self.moving_variance * self.momentum
+                + variance * (1.0 - self.momentum)
+            )
         else:
             moving_mean = ops.reshape(self.moving_mean, broadcast_shape)
             moving_variance = ops.reshape(self.moving_variance, broadcast_shape)


Test patch:
diff --git a/keras_core/layers/normalization/batch_normalization_test.py b/keras_core/layers/normalization/batch_normalization_test.py
--- a/keras_core/layers/normalization/batch_normalization_test.py
+++ b/keras_core/layers/normalization/batch_normalization_test.py
@@ -61,10 +61,23 @@ def test_bn_basics(self):
     @parameterized.product(
         axis=(-1, 1),
         input_shape=((5, 2, 3), (5, 3, 3, 2)),
+        moving_mean_initializer=("zeros", "ones"),
+        moving_variance_initializer=("zeros", "ones"),
     )
-    def test_correctness(self, axis, input_shape):
+    def test_correctness(
+        self,
+        axis,
+        input_shape,
+        moving_mean_initializer,
+        moving_variance_initializer,
+    ):
         # Training
-        layer = layers.BatchNormalization(axis=axis, momentum=0)
+        layer = layers.BatchNormalization(
+            axis=axis,
+            momentum=0,
+            moving_mean_initializer=moving_mean_initializer,
+            moving_variance_initializer=moving_variance_initializer,
+        )
         # Random data centered on 5.0, variance 10.0
         x = np.random.normal(loc=5.0, scale=10.0, size=input_shape)
         out = x
@@ -83,12 +96,19 @@ def test_correctness(self, axis, input_shape):
         reduction_axes = tuple(reduction_axes)
         self.assertAllClose(np.mean(out, axis=reduction_axes), 0.0, atol=1e-3)
         self.assertAllClose(np.std(out, axis=reduction_axes), 1.0, atol=1e-3)
+        self.assertAllClose(layer.moving_mean, 0.0, atol=1e-3)
+        self.assertAllClose(layer.moving_variance, 1.0, atol=1e-3)
 
-        # Inference
+        # Inference done before training shouldn't match.
         inference_out = layer(x, training=False)
         training_out = layer(x, training=True)
         self.assertNotAllClose(inference_out, training_out)
 
+        # Since momentum is zero, inference after training should match.
+        training_out = layer(x, training=True)
+        inference_out = layer(x, training=False)
+        self.assertAllClose(inference_out, training_out)
+
     def test_trainable_behavior(self):
         layer = layers.BatchNormalization(axis=-1, momentum=0.8, epsilon=1e-7)
         layer.build((1, 4, 4, 3))