Fix Layer normalization issue with scalar mean & variance (#20626)

* Fix normalization issue with scalar mean & variance

* Add unit test for normalization with scalar mean and variance

keras/src/layers/preprocessing/normalization.py

@@ -190,8 +190,8 @@ def build(self, input_shape):
             # with proper broadcast shape for use during call.
             mean = ops.convert_to_tensor(self.input_mean)
             variance = ops.convert_to_tensor(self.input_variance)
-            mean = ops.reshape(mean, self._broadcast_shape)
-            variance = ops.reshape(variance, self._broadcast_shape)
+            mean = ops.broadcast_to(mean, self._broadcast_shape)
+            variance = ops.broadcast_to(variance, self._broadcast_shape)
             self.mean = ops.cast(mean, dtype=self.compute_dtype)
             self.variance = ops.cast(variance, dtype=self.compute_dtype)
             self.built = True

keras/src/layers/preprocessing/normalization_test.py

@@ -164,3 +164,8 @@ def test_tf_data_compatibility(self):
         )
         for output in ds.map(layer).take(1):
             output.numpy()
+
+    def test_normalization_with_scalar_mean_var(self):
+        input_data = np.array([[1,2,3]], dtype='float32')
+        layer = layers.Normalization(mean=3., variance=2.)
+        layer(input_data)