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<int>{0}, keepdims, s);
+    auto out = norm(flatten(a, s), std::vector<int>{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/