Reject a negative eps in the normalization layers

python/mlx/nn/layers/normalization.py

@@ -47,6 +47,8 @@ def __init__(
         affine: bool = False,
     ):
         super().__init__()
+        if eps <= 0.0:
+            raise ValueError(f"[InstanceNorm] 'eps' must be positive but got {eps}.")
         if affine:
             self.weight = mx.ones((dims,))
             self.bias = mx.zeros((dims,))
@@ -101,6 +103,8 @@ def __init__(
         self, dims: int, eps: float = 1e-5, affine: bool = True, bias: bool = True
     ):
         super().__init__()
+        if eps <= 0.0:
+            raise ValueError(f"[LayerNorm] 'eps' must be positive but got {eps}.")
         if affine:
             self.weight = mx.ones((dims,))
             if bias:
@@ -141,6 +145,8 @@ class RMSNorm(Module):
 
     def __init__(self, dims: int, eps: float = 1e-5):
         super().__init__()
+        if eps <= 0.0:
+            raise ValueError(f"[RMSNorm] 'eps' must be positive but got {eps}.")
         self.weight = mx.ones((dims,))
         self.eps = eps
 
@@ -191,6 +197,8 @@ def __init__(
         pytorch_compatible: bool = False,
     ):
         super().__init__()
+        if eps <= 0.0:
+            raise ValueError(f"[GroupNorm] 'eps' must be positive but got {eps}.")
         if num_groups <= 0:
             raise ValueError(
                 f"The number of groups ({num_groups}) must be a positive integer."
@@ -309,6 +317,8 @@ def __init__(
         track_running_stats: bool = True,
     ):
         super().__init__()
+        if eps <= 0.0:
+            raise ValueError(f"[BatchNorm] 'eps' must be positive but got {eps}.")
 
         self.num_features = num_features
         self.eps = eps

python/tests/test_nn.py

@@ -410,6 +410,28 @@ def test_bilinear(self):
         outputs = layer(inputs1, inputs2)
         self.assertEqual(outputs.shape, (10, 6))
 
+    def test_norm_eps_validation(self):
+        # eps is added under a square root. A negative one makes rsqrt take the
+        # root of a negative number, so the layer emits NaN for whichever
+        # elements have a small enough variance, which is only a partial NaN and
+        # easy to miss. Zero is rejected too: it leaves rsqrt(0) for any input
+        # whose variance is zero, which is a whole NaN row. This matches the eps
+        # guards the optimizers already carry.
+        builders = (
+            ("LayerNorm", lambda eps: nn.LayerNorm(16, eps=eps)),
+            ("RMSNorm", lambda eps: nn.RMSNorm(16, eps=eps)),
+            ("GroupNorm", lambda eps: nn.GroupNorm(4, 16, eps=eps)),
+            ("InstanceNorm", lambda eps: nn.InstanceNorm(16, eps=eps)),
+            ("BatchNorm", lambda eps: nn.BatchNorm(16, eps=eps)),
+        )
+        for name, build in builders:
+            for eps in (-1.0, -1e-30, 0.0):
+                with self.assertRaisesRegex(ValueError, "must be positive"):
+                    build(eps)
+            # Anything positive still constructs, including a very small eps.
+            for eps in (1e-30, 1e-5, 1.0):
+                build(eps)
+
     def test_group_norm(self):
         x = mx.arange(100, dtype=mx.float32)
         x = x.reshape(1, 10, 10, 1)