[Fix] Mixed Precision (float16) numerical instability in GroupNormalization with small epsilon

keras/src/layers/normalization/group_normalization.py

@@ -1,3 +1,7 @@
+import warnings
+
+import ml_dtypes
+
 from keras.src import backend
 from keras.src import constraints
 from keras.src import initializers
@@ -93,6 +97,22 @@ def __init__(
         self.beta_constraint = constraints.get(beta_constraint)
         self.gamma_constraint = constraints.get(gamma_constraint)
 
+        compute_dtype = backend.standardize_dtype(self.compute_dtype)
+        try:
+            finfo = ml_dtypes.finfo(compute_dtype)
+            if self.epsilon != 0 and self.epsilon < float(finfo.eps):
+                warnings.warn(
+                    "The configured `epsilon` is smaller than what can be "
+                    f"represented in the layer compute dtype "
+                    f"({compute_dtype}); "
+                    "it may be rounded to 0 under autocast. Consider "
+                    "increasing `epsilon` or setting `autocast=False` "
+                    "for this layer.",
+                    stacklevel=2,
+                )
+        except Exception:
+            pass
+
     def build(self, input_shape):
         dim = input_shape[self.axis]
 
@@ -151,7 +171,8 @@ def call(self, inputs):
         normalized_inputs = self._apply_normalization(
             reshaped_inputs, inputs.shape
         )
-        return ops.reshape(normalized_inputs, ops.shape(inputs))
+        outputs = ops.reshape(normalized_inputs, ops.shape(inputs))
+        return ops.cast(outputs, self.compute_dtype)
 
     def _reshape_into_groups(self, inputs):
         input_shape = ops.shape(inputs)

keras/src/layers/normalization/group_normalization_test.py

@@ -1,4 +1,7 @@
+import warnings
+
 import numpy as np
+import pytest
 
 from keras.src import constraints
 from keras.src import layers
@@ -175,3 +178,26 @@ def test_broadcasting_2d_channels_first(self):
             ),
             atol=1e-3,
         )
+
+    def test_warns_when_epsilon_too_small_for_compute_dtype(self):
+        with pytest.warns(UserWarning, match="epsilon"):
+            layers.GroupNormalization(
+                groups=2,
+                axis=-1,
+                scale=False,
+                center=False,
+                epsilon=1e-4,
+                dtype="mixed_bfloat16",
+            )
+
+        with warnings.catch_warnings(record=True) as w:
+            warnings.simplefilter("always")
+            layers.GroupNormalization(
+                groups=2,
+                axis=-1,
+                scale=False,
+                center=False,
+                epsilon=1e-3,
+                dtype="mixed_bfloat16",
+            )
+            self.assertFalse(any("epsilon" in str(x.message) for x in w))