From 9e9b35a4b7fbe7b987cafe2bde482dbcbed86739 Mon Sep 17 00:00:00 2001 From: devteamaegis Date: Mon, 10 Aug 2026 16:04:15 -0400 Subject: [PATCH] fix(linalg): keep the input rank in norm when axis and ord are both None With no axis and no ord, norm flattens the input before reducing, so keepdims restored the rank of the flattened array instead of the original one. A (2, 3, 4) input returned shape (1,) instead of (1, 1, 1). Reshape the result back to the input rank. Every other path was already correct: an explicit ord, an explicit axis, and the other reductions (sum, mean, max, var) all keep the rank. The existing test already covered this case but compared with np.allclose, which broadcasts (1,) against (1, 1) and passes. Assert the shape alongside the values so the blind spot is closed. --- mlx/linalg.cpp | 9 ++++++++- python/tests/test_linalg.py | 3 +++ 2 files changed, 11 insertions(+), 1 deletion(-) diff --git a/mlx/linalg.cpp b/mlx/linalg.cpp index 159d2468bc..edfb9cb79a 100644 --- a/mlx/linalg.cpp +++ b/mlx/linalg.cpp @@ -166,7 +166,14 @@ array norm( bool keepdims /* = false */, StreamOrDevice s /* = {} */) { if (!axis) { - return norm(flatten(a, s), std::vector{0}, keepdims, s); + auto out = norm(flatten(a, s), std::vector{0}, keepdims, s); + if (keepdims) { + // The flatten above collapses the input to one dimension, so keepdims + // has to restore the rank of the original array rather than the + // flattened one. + out = reshape(out, Shape(a.ndim(), 1), s); + } + return out; } if (axis.value().size() > 2) { diff --git a/python/tests/test_linalg.py b/python/tests/test_linalg.py index 9b3859976f..52e02e4d9a 100644 --- a/python/tests/test_linalg.py +++ b/python/tests/test_linalg.py @@ -38,6 +38,7 @@ def test_norm(self): with self.subTest( shape=shape, ord=o, axis=axis, keepdims=keepdims ): + self.assertEqual(out_mx.shape, out_np.shape) self.assertTrue( np.allclose(out_np, out_mx, atol=1e-5, rtol=1e-6) ) @@ -51,6 +52,7 @@ def test_norm(self): out_np = np.linalg.norm(x_np, ord=o, keepdims=keepdims) out_mx = mx.linalg.norm(x_mx, ord=o, keepdims=keepdims) with self.subTest(shape=shape, ord=o, keepdims=keepdims): + self.assertEqual(out_mx.shape, out_np.shape) self.assertTrue( np.allclose(out_np, out_mx, atol=1e-5, rtol=1e-6) ) @@ -63,6 +65,7 @@ def test_norm(self): out_np = np.linalg.norm(x_np, keepdims=keepdims) out_mx = mx.linalg.norm(x_mx, keepdims=keepdims) with self.subTest(shape=shape, keepdims=keepdims): + self.assertEqual(out_mx.shape, out_np.shape) self.assertTrue(np.allclose(out_np, out_mx, atol=1e-5, rtol=1e-6)) # tests for negative indexing: -1/1/inf/-inf/