added validation checks in Group, Layer, Batch Normalization layers (#…

…20246)

* added validation checks in Group, Layer, Batch Normalization layers for the compute_output_shape function

* Update batch_normalization.py

* Update group_normalization.py

* Update layer_normalization.py

---------

Co-authored-by: François Chollet <[email protected]>

keras/src/layers/normalization/batch_normalization.py

@@ -219,6 +219,18 @@ def build(self, input_shape):
         self.built = True
 
     def compute_output_shape(self, input_shape):
+        if isinstance(self.axis, int):
+            axes = [self.axis]
+        else:
+            axes = self.axis
+
+        for axis in axes:
+            if axis >= len(input_shape) or axis < -len(input_shape):
+                raise ValueError(
+                    f"Axis {axis} is out of bounds for "
+                    f"input shape {input_shape}. "
+                    f"Received: axis={self.axis}"
+                )
         return input_shape
 
     def call(self, inputs, training=None, mask=None):

keras/src/layers/normalization/group_normalization.py

@@ -199,6 +199,18 @@ def _create_broadcast_shape(self, input_shape):
         return broadcast_shape
 
     def compute_output_shape(self, input_shape):
+        if isinstance(self.axis, int):
+            axes = [self.axis]
+        else:
+            axes = self.axis
+
+        for axis in axes:
+            if axis >= len(input_shape) or axis < -len(input_shape):
+                raise ValueError(
+                    f"Axis {axis} is out of bounds for "
+                    f"input shape {input_shape}. "
+                    f"Received: axis={self.axis}"
+                )
         return input_shape
 
     def get_config(self):

keras/src/layers/normalization/layer_normalization.py

@@ -117,7 +117,7 @@ def __init__(
         gamma_regularizer=None,
         beta_constraint=None,
         gamma_constraint=None,
-        **kwargs
+        **kwargs,
     ):
         super().__init__(**kwargs)
         if isinstance(axis, (list, tuple)):
@@ -235,6 +235,18 @@ def _broadcast(v):
         return ops.cast(outputs, input_dtype)
 
     def compute_output_shape(self, input_shape):
+        if isinstance(self.axis, int):
+            axes = [self.axis]
+        else:
+            axes = self.axis
+
+        for axis in axes:
+            if axis >= len(input_shape) or axis < -len(input_shape):
+                raise ValueError(
+                    f"Axis {axis} is out of bounds for "
+                    f"input shape {input_shape}. "
+                    f"Received: axis={self.axis}"
+                )
         return input_shape
 
     def get_config(self):