Fixes an axis mismatch bug in matrix norm for case -1 and 1

[danlee2002]

Proposed changes

The following PR aims to fix a matrix norm calculation bug caused by similar reason #3756 for when we are consider the $\pm1$ norm. Since the matrix norm is calculated as a series of reduction. When keepdims is false and negatives axes are passed in, the method fails to adjust the col_axis accordingly to account for the loss in dimension. This results in either an index out of bound error or the wrong index being reduced along. The following PR normalizes the axes accordingly to account for this reduction.

Reproduction

To reproduce the bug

import mlx.core as mx
import numpy as np
import mlx.core.linalg as la

a_mx = mx.arange(9).reshape(3,3)
a_np = np.arange(9).reshape(3,3)
print(np.linalg.norm(a_np, ord = -1, axis = (-1, -2))) #5.0
print(la.norm(a_mx, ord = -1, axis = [-1, -2])) # throws exception

a_mx = mx.arange(18).reshape(2,3,3)
a_np = np.arange(18).reshape(2,3,3)
print(np.linalg.norm(a_np, ord = -1, axis = (-1, -2))) # [21., 48.]
print(la.norm(a_mx, ord = -1, axis = [-1, -2])) # array([3, 12, 21], dtype=float32)

Testing

• Existing tests were ran
• Existing test suite for negative axis reduction was extended to account for $\pm 1$ norm to account for all possible negative axes combination for (3, 3), (2, 3, 3), (2, 3, 3, 3) matrices

Checklist

Put an x in the boxes that apply.

• 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)

[zcbenz]

Thanks!