From d1df6d9fda5c78c628712bccc0044ef4e005d3ba Mon Sep 17 00:00:00 2001 From: ayaangazali Date: Sun, 16 Aug 2026 13:52:35 -0700 Subject: [PATCH 1/2] Reject a negative eps in the normalization layers eps is added under a square root, so a negative one makes rsqrt take the root of a negative number and the layer returns NaN for whichever elements have a small enough variance. That is a partial NaN rather than an obvious all NaN output. Zero is still allowed. --- python/mlx/nn/layers/normalization.py | 12 ++++++++++++ python/tests/test_nn.py | 20 ++++++++++++++++++++ 2 files changed, 32 insertions(+) diff --git a/python/mlx/nn/layers/normalization.py b/python/mlx/nn/layers/normalization.py index e79440dce3..285d0eaa61 100644 --- a/python/mlx/nn/layers/normalization.py +++ b/python/mlx/nn/layers/normalization.py @@ -52,6 +52,10 @@ def __init__( self.bias = mx.zeros((dims,)) self.dims = dims self.eps = eps + if eps < 0.0: + raise ValueError( + f"[InstanceNorm] 'eps' must be non-negative but got {eps}." + ) def _extra_repr(self): return f"{self.dims}, eps={self.eps}, affine={'weight' in self}" @@ -106,6 +110,8 @@ def __init__( if bias: self.bias = mx.zeros((dims,)) self.eps = eps + if eps < 0.0: + raise ValueError(f"[LayerNorm] 'eps' must be non-negative but got {eps}.") self.dims = dims def _extra_repr(self): @@ -143,6 +149,8 @@ def __init__(self, dims: int, eps: float = 1e-5): super().__init__() self.weight = mx.ones((dims,)) self.eps = eps + if eps < 0.0: + raise ValueError(f"[RMSNorm] 'eps' must be non-negative but got {eps}.") def _extra_repr(self): return f"{self.weight.shape[0]}, eps={self.eps}" @@ -210,6 +218,8 @@ def __init__( self.num_groups = num_groups self.dims = dims self.eps = eps + if eps < 0.0: + raise ValueError(f"[GroupNorm] 'eps' must be non-negative but got {eps}.") self.pytorch_compatible = pytorch_compatible def _extra_repr(self): @@ -312,6 +322,8 @@ def __init__( self.num_features = num_features self.eps = eps + if eps < 0.0: + raise ValueError(f"[BatchNorm] 'eps' must be non-negative but got {eps}.") self.momentum = momentum self.track_running_stats = track_running_stats diff --git a/python/tests/test_nn.py b/python/tests/test_nn.py index c5e6db94a7..762c4af50e 100644 --- a/python/tests/test_nn.py +++ b/python/tests/test_nn.py @@ -410,6 +410,26 @@ 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, so a negative one makes rsqrt take + # the root of a negative number and the layer emits NaN for whichever + # elements happen to have a small enough variance. That is only a + # partial NaN, so it is easy to miss. Reject it up front instead. + 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): + with self.assertRaisesRegex(ValueError, "non-negative"): + build(eps) + # Zero is a legitimate choice and still constructs. + for eps in (0.0, 1e-5): + build(eps) + def test_group_norm(self): x = mx.arange(100, dtype=mx.float32) x = x.reshape(1, 10, 10, 1) From 91f2876d05d0d0be363961d076d73538b718846d Mon Sep 17 00:00:00 2001 From: ayaangazali Date: Mon, 17 Aug 2026 23:53:16 -0700 Subject: [PATCH 2/2] Reject eps == 0 too, and validate before assigning members Zero leaves rsqrt(0) for any input whose variance is zero, so require a positive eps rather than merely non-negative, matching the eps guards in RMSprop, Adagrad and AdaDelta. Guards now run right after super().__init__() like GroupNorm's do. --- python/mlx/nn/layers/normalization.py | 22 ++++++++++------------ python/tests/test_nn.py | 18 ++++++++++-------- 2 files changed, 20 insertions(+), 20 deletions(-) diff --git a/python/mlx/nn/layers/normalization.py b/python/mlx/nn/layers/normalization.py index 285d0eaa61..a5a8293fdb 100644 --- a/python/mlx/nn/layers/normalization.py +++ b/python/mlx/nn/layers/normalization.py @@ -47,15 +47,13 @@ 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,)) self.dims = dims self.eps = eps - if eps < 0.0: - raise ValueError( - f"[InstanceNorm] 'eps' must be non-negative but got {eps}." - ) def _extra_repr(self): return f"{self.dims}, eps={self.eps}, affine={'weight' in self}" @@ -105,13 +103,13 @@ 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: self.bias = mx.zeros((dims,)) self.eps = eps - if eps < 0.0: - raise ValueError(f"[LayerNorm] 'eps' must be non-negative but got {eps}.") self.dims = dims def _extra_repr(self): @@ -147,10 +145,10 @@ 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 - if eps < 0.0: - raise ValueError(f"[RMSNorm] 'eps' must be non-negative but got {eps}.") def _extra_repr(self): return f"{self.weight.shape[0]}, eps={self.eps}" @@ -199,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." @@ -218,8 +218,6 @@ def __init__( self.num_groups = num_groups self.dims = dims self.eps = eps - if eps < 0.0: - raise ValueError(f"[GroupNorm] 'eps' must be non-negative but got {eps}.") self.pytorch_compatible = pytorch_compatible def _extra_repr(self): @@ -319,11 +317,11 @@ 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 - if eps < 0.0: - raise ValueError(f"[BatchNorm] 'eps' must be non-negative but got {eps}.") self.momentum = momentum self.track_running_stats = track_running_stats diff --git a/python/tests/test_nn.py b/python/tests/test_nn.py index 762c4af50e..46242d192c 100644 --- a/python/tests/test_nn.py +++ b/python/tests/test_nn.py @@ -411,10 +411,12 @@ def test_bilinear(self): self.assertEqual(outputs.shape, (10, 6)) def test_norm_eps_validation(self): - # eps is added under a square root, so a negative one makes rsqrt take - # the root of a negative number and the layer emits NaN for whichever - # elements happen to have a small enough variance. That is only a - # partial NaN, so it is easy to miss. Reject it up front instead. + # 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)), @@ -423,11 +425,11 @@ def test_norm_eps_validation(self): ("BatchNorm", lambda eps: nn.BatchNorm(16, eps=eps)), ) for name, build in builders: - for eps in (-1.0, -1e-30): - with self.assertRaisesRegex(ValueError, "non-negative"): + for eps in (-1.0, -1e-30, 0.0): + with self.assertRaisesRegex(ValueError, "must be positive"): build(eps) - # Zero is a legitimate choice and still constructs. - for eps in (0.0, 1e-5): + # 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):