diff --git a/FE b/FE new file mode 100644 index 00000000000..e69de29bb2d diff --git a/keras/src/layers/normalization/batch_normalization.py b/keras/src/layers/normalization/batch_normalization.py index 8b0160344ee..67f12cab956 100644 --- a/keras/src/layers/normalization/batch_normalization.py +++ b/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): diff --git a/keras/src/layers/normalization/group_normalization.py b/keras/src/layers/normalization/group_normalization.py index f70fb69f3ed..c547c99a6b9 100644 --- a/keras/src/layers/normalization/group_normalization.py +++ b/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): diff --git a/keras/src/layers/normalization/layer_normalization.py b/keras/src/layers/normalization/layer_normalization.py index 73a8956fd8f..6df9fd1df89 100644 --- a/keras/src/layers/normalization/layer_normalization.py +++ b/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):