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

…eras-team#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 <francois.chollet@gmail.com>

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):