fix(linalg): keep the input rank in norm with keepdims and no axis

[devteamaegis]

Proposed changes

What's broken. mx.linalg.norm(x, keepdims=True) collapses the result to a single dimension instead of keeping the input rank, when both axis and ord are left at their defaults:

import mlx.core as mx
mx.linalg.norm(mx.ones((2, 3, 4)), keepdims=True).shape   # (1,)      -- expected (1, 1, 1)

input	mlx	numpy
(4,)	(1,)	(1,)
(2, 3)	(1,)	(1, 1)
(2, 3, 4)	(1,)	(1, 1, 1)
(2, 3, 4, 5)	(1,)	(1, 1, 1, 1)

A 0-d input is wrong in the other direction: mlx gives (1,) where NumPy gives ().

Every neighbouring path is already correct, so this is specifically the default-argument path:

• an explicit ord works — ord=2, 1, inf, "fro", "nuc" all give (1, 1) for a (2, 3) input
• an explicit axis works — axis=0, 1, (0, 1) all match NumPy
• the other reductions work — sum, mean, max, var with keepdims=True all give (1, 1)

Why. With no axis, norm flattens the input and reduces axis 0, so keepdims preserves the rank of the flattened array (always 1) rather than the original. The two sibling overloads that take an ord don't flatten — they build the axis list over the original array — which is why they were unaffected.

The fix. Reshape the result back to the input rank when keepdims is set. The computed value is untouched; only the shape changes. This also fixes the 0-d case, since a rank-0 input reshapes back to a scalar.

The test. The existing test_norm already exercised this exact case, but compared with np.allclose, which broadcasts (1,) against (1, 1) and passes — so the shape bug was invisible. I added assertEqual(out_mx.shape, out_np.shape) alongside the value check in all three test_norm loops, which closes the blind spot generally rather than only for this bug.

Fails before, passes after:

# before
SUBFAILED(shape=(2, 3), ord=None, keepdims=True) ... - (1,) + (1, 1)
SUBFAILED(shape=(2, 3), keepdims=True)
SUBFAILED(shape=(2, 3, 3), keepdims=True)
3 failed, 1 passed, 17 deselected, 247 subtests passed

# after
18 passed, 266 subtests passed

Only those three subtests failed beforehand, so the added shape assertions do not flag anything else.

Benchmark. The added reshape is metadata-only and runs only when keepdims is true. CPU, M4, best of 3, using benchmarks/python/time_utils.py:

case	before	after
norm((1_000_000,)) keepdims=False	0.1518 ms	0.1490 ms
norm((1_000_000,)) keepdims=True	0.1538 ms	0.1477 ms
norm((1024, 1024)) keepdims=False	0.1578 ms	0.1564 ms
norm((1024, 1024)) keepdims=True	0.1545 ms	0.1563 ms
norm((128, 128, 64)) keepdims=False	0.1554 ms	0.1558 ms
norm((128, 128, 64)) keepdims=True	0.1550 ms	0.1569 ms
norm((1024, 1024), ord="fro", keepdims=True) (control)	0.1572 ms	0.1556 ms

All within ~2% of each other and of the untouched control, i.e. indistinguishable from noise.

Checklist

• I have read the CONTRIBUTING document
• I have run pre-commit run --all-files to format my code / installed pre-commit prior to committing changes
• I have added tests that prove my fix is effective or that my feature works
• I have updated the necessary documentation (if needed) — no API change; the docstring already describes the NumPy behaviour this restores