Repository: keras-team/keras (python)

Issue:
Normalization layer fails to broadcast scalar mean and variance
Normalization layer errors out for scalar mean and variance arguments. The documentation states that these arguments will be broadcast to the necessary size and the examples include scalars, but instead the source code seems to only reshape them (see [here](https://github.com/keras-team/keras/blob/v3.7.0/keras/src/layers/preprocessing/normalization.py#L193)).

I'm using Python 3.12.1 and Keras 3.6.0.

```
>>> in1 = keras.Input(shape=(16, 16, 3))
>>> normLayer = keras.layers.Normalization(mean=0.1, variance=0.05)
>>> out1 = normLayer(in1)
I tensorflow/core/framework/local_rendezvous.cc:405] Local rendezvous is aborting with status: INVALID_ARGUMENT: Input to reshape is a tensor with 1 values, but the requested shape has 3
```

Pull request description:
Fix Layer normalization issue with scalar mean & variance
Fix Layer normalization issue with scalar mean & variance.

With scalar mean instead of broadcasting reshaping applied currently. Changed the same.

Fixes #20624.

Patch:
diff --git a/keras/src/backend/jax/distribution_lib.py b/keras/src/backend/jax/distribution_lib.py
index d086c650e..edb1dc118 100644
--- a/keras/src/backend/jax/distribution_lib.py
+++ b/keras/src/backend/jax/distribution_lib.py
@@ -1,10 +1,4 @@
-"""!!!DO NOT USE!!!
-
-Distribution related class for JAX backend.
-
-This is just a prototype and we might want to unify it
-with other backends in the future.
-"""
+"""Utilities for distribution strategy with JAX backend."""
 
 import jax
 import numpy as np
diff --git a/keras/src/backend/tensorflow/random.py b/keras/src/backend/tensorflow/random.py
index 4bd2162fc..e807b0de9 100644
--- a/keras/src/backend/tensorflow/random.py
+++ b/keras/src/backend/tensorflow/random.py
@@ -94,15 +94,11 @@ def dropout(inputs, rate, noise_shape=None, seed=None):
 
 
 def shuffle(x, axis=0, seed=None):
-    from keras.src.backend.tensorflow.numpy import swapaxes
-
     seed = _cast_seed(draw_seed(seed))
-    if axis == 0:
-        return tf.random.experimental.stateless_shuffle(x, seed=seed)
-    x = swapaxes(x, axis1=0, axis2=axis)
-    x = tf.random.experimental.stateless_shuffle(x, seed=seed)
-    x = swapaxes(x, axis1=0, axis2=axis)
-    return x
+    indices = tf.argsort(
+        tf.random.stateless_uniform(shape=[tf.shape(x)[axis]], seed=seed)
+    )
+    return tf.gather(x, indices, axis=axis)
 
 
 def gamma(shape, alpha, dtype=None, seed=None):
diff --git a/keras/src/distribution/distribution_lib.py b/keras/src/distribution/distribution_lib.py
index f161d575d..b4736426a 100644
--- a/keras/src/distribution/distribution_lib.py
+++ b/keras/src/distribution/distribution_lib.py
@@ -1,9 +1,7 @@
-"""Unified high level distribution APIs across backends.
+"""Unified high-level distribution APIs across backends.
 
-!!!DO NOT USE!!! Currently under development and APIs are not final.
-
-Currently only the JAX backend has been implemented. The TensorFlow backend
-will be implemented in the future (via tf.dtensor API).
+Currently only the JAX backend is supported. The TensorFlow backend
+will be supported in the future (via tf.dtensor API).
 """
 
 import collections
diff --git a/keras/src/layers/preprocessing/image_preprocessing/mix_up.py b/keras/src/layers/preprocessing/image_preprocessing/mix_up.py
index 7c75acd0d..ba0b92280 100644
--- a/keras/src/layers/preprocessing/image_preprocessing/mix_up.py
+++ b/keras/src/layers/preprocessing/image_preprocessing/mix_up.py
@@ -1,4 +1,3 @@
-import keras.src.random.random
 from keras.src.api_export import keras_export
 from keras.src.layers.preprocessing.image_preprocessing.base_image_preprocessing_layer import (  # noqa: E501
     BaseImagePreprocessingLayer,
@@ -36,7 +35,7 @@ class MixUp(BaseImagePreprocessingLayer):
     """
 
     def __init__(self, alpha=0.2, data_format=None, seed=None, **kwargs):
-        super().__init__(data_format=None, **kwargs)
+        super().__init__(data_format=data_format, **kwargs)
         self.alpha = alpha
         self.seed = seed
         self.generator = SeedGenerator(seed)
@@ -54,13 +53,16 @@ class MixUp(BaseImagePreprocessingLayer):
         else:
             batch_size = self.backend.shape(images)[0]
 
+        if seed is None:
+            seed = self._get_seed_generator(self.backend._backend)
+
         permutation_order = self.backend.random.shuffle(
             self.backend.numpy.arange(0, batch_size, dtype="int64"),
-            seed=self.generator,
+            seed=seed,
         )
 
-        mix_weight = keras.src.random.random.beta(
-            (1,), self.alpha, self.alpha, seed=self.generator
+        mix_weight = self.backend.random.beta(
+            (1,), self.alpha, self.alpha, seed=seed
         )
         return {
             "mix_weight": mix_weight,
diff --git a/keras/src/layers/preprocessing/normalization.py b/keras/src/layers/preprocessing/normalization.py
index e39dae73d..54ed025c9 100644
--- a/keras/src/layers/preprocessing/normalization.py
+++ b/keras/src/layers/preprocessing/normalization.py
@@ -190,8 +190,8 @@ class Normalization(TFDataLayer):
             # 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


Test patch:
diff --git a/keras/src/layers/preprocessing/normalization_test.py b/keras/src/layers/preprocessing/normalization_test.py
index 4278272a5..473f5daf6 100644
--- a/keras/src/layers/preprocessing/normalization_test.py
+++ b/keras/src/layers/preprocessing/normalization_test.py
@@ -164,3 +164,8 @@ class NormalizationTest(testing.TestCase):
         )
         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)