diff --git a/python/mlx/nn/layers/normalization.py b/python/mlx/nn/layers/normalization.py index e79440dce3..69b15e54da 100644 --- a/python/mlx/nn/layers/normalization.py +++ b/python/mlx/nn/layers/normalization.py @@ -62,12 +62,19 @@ def __call__(self, x: mx.array) -> mx.array: f"InstanceNorm expects inputs with at least 3 dimensions" f" (N, ..., C) but the input has {x.ndim} dimensions." ) - reduction_axes = tuple(range(1, x.ndim - 1)) - # Compute stats - mean = mx.mean(x, axis=reduction_axes, keepdims=True) - var = mx.var(x, axis=reduction_axes, keepdims=True) - # Normalize - x = (x - mean) * mx.rsqrt(var + self.eps) + batch_size, features = x.shape[0], x.shape[-1] + spatial_shape = x.shape[1:-1] + channels_first = mx.transpose(x, (0, x.ndim - 1, *range(1, x.ndim - 1))) + x = mx.fast.layer_norm( + channels_first.reshape(batch_size, features, -1), + None, + None, + self.eps, + ) + x = mx.transpose( + x.reshape(batch_size, features, *spatial_shape), + (0, *range(2, len(spatial_shape) + 2), 1), + ) # Scale and shift if necessary return (self.weight * x + self.bias) if "weight" in self else x diff --git a/python/tests/test_nn.py b/python/tests/test_nn.py index c5e6db94a7..67d555b4a2 100644 --- a/python/tests/test_nn.py +++ b/python/tests/test_nn.py @@ -671,6 +671,21 @@ def test_instance_norm(self): ] self.assertTrue(x.shape == y.shape) self.assertTrue(np.allclose(y, expected_y, atol=1e-5)) + # Reduced-precision statistics must not overflow for finite feature maps. + checkerboard = np.indices((4, 4, 4)).sum(axis=0) % 2 + x = mx.array( + np.stack( + [ + np.where(checkerboard, -512, 512), + np.where(checkerboard, -256, 256), + ], + axis=-1, + ).astype(np.float16) + )[None] + y = nn.InstanceNorm(dims=2)(x) + self.assertEqual(y.dtype, mx.float16) + self.assertTrue(mx.allclose(y.min(), mx.array(-1.0, dtype=mx.float16))) + self.assertTrue(mx.allclose(y.max(), mx.array(1.0, dtype=mx.float16))) # Test repr self.assertTrue(str(inorm) == "InstanceNorm(3, eps=1e-05, affine=False)") # Raise for inputs without spatial dimensions