diff --git a/python/mlx/nn/layers/normalization.py b/python/mlx/nn/layers/normalization.py index e79440dce3..a5a8293fdb 100644 --- a/python/mlx/nn/layers/normalization.py +++ b/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 diff --git a/python/tests/test_nn.py b/python/tests/test_nn.py index c5e6db94a7..46242d192c 100644 --- a/python/tests/test_nn.py +++ b/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)