diff --git a/mlx/linalg.cpp b/mlx/linalg.cpp index bdf05fb80b..75a662b1e6 100644 --- a/mlx/linalg.cpp +++ b/mlx/linalg.cpp @@ -84,8 +84,8 @@ inline array matrix_norm( bool keepdims, StreamOrDevice s) { auto dtype = at_least_float(a.dtype()); - auto row_axis = axis[0]; - auto col_axis = axis[1]; + int row_axis = (axis[0] < 0) ? axis[0] + a.ndim() : axis[0]; + int col_axis = (axis[1] < 0) ? axis[1] + a.ndim() : axis[1]; if (ord == -1.0) { col_axis -= (!keepdims && col_axis > row_axis && col_axis > 0); return astype( @@ -99,24 +99,18 @@ inline array matrix_norm( dtype, s); } else if (ord == std::numeric_limits::infinity()) { - row_axis = (axis[0] < 0) ? axis[0] + a.ndim() : axis[0]; - col_axis = (axis[1] < 0) ? axis[1] + a.ndim() : axis[1]; row_axis -= (!keepdims && row_axis > col_axis && row_axis > 0); return astype( max(sum(abs(a, s), col_axis, keepdims, s), row_axis, keepdims, s), dtype, s); } else if (ord == -std::numeric_limits::infinity()) { - row_axis = (axis[0] < 0) ? axis[0] + a.ndim() : axis[0]; - col_axis = (axis[1] < 0) ? axis[1] + a.ndim() : axis[1]; row_axis -= (!keepdims && row_axis > col_axis && row_axis > 0); return astype( min(sum(abs(a, s), col_axis, keepdims, s), row_axis, keepdims, s), dtype, s); } else if (ord == 2.0 || ord == -2.0) { - row_axis = (axis[0] < 0) ? axis[0] + a.ndim() : axis[0]; - col_axis = (axis[1] < 0) ? axis[1] + a.ndim() : axis[1]; auto a_matrix = (row_axis > col_axis) ? moveaxis(moveaxis(a, row_axis, -1, s), col_axis, -1, s) : moveaxis(moveaxis(a, col_axis, -1, s), row_axis, -2, s); diff --git a/python/tests/test_linalg.py b/python/tests/test_linalg.py index afdf75d7c8..39f0a3b891 100644 --- a/python/tests/test_linalg.py +++ b/python/tests/test_linalg.py @@ -65,8 +65,8 @@ def test_norm(self): with self.subTest(shape=shape, keepdims=keepdims): self.assertTrue(np.allclose(out_np, out_mx, atol=1e-5, rtol=1e-6)) - # neg/pos inf norm test - norms = [-float("inf"), float("inf")] + # tests for negative indexing: -1/1/inf/-inf/ + norms = [-1, 1, -float("inf"), float("inf")] for shape in [(3, 3), (2, 3, 3), (2, 3, 3, 3)]: x_mx = mx.arange(1, math.prod(shape) + 1, dtype=mx.float32).reshape(shape) x_np = np.arange(1, math.prod(shape) + 1, dtype=np.float32).reshape(shape) @@ -76,7 +76,7 @@ def test_norm(self): for axes in neg_axes: out_np = np.linalg.norm( x_np, - ord=np.inf if ord == float("inf") else -np.inf, + ord=ord, axis=tuple(axes), ) out_mx = mx.linalg.norm(x_mx, ord=ord, axis=axes) diff --git a/tests/linalg_tests.cpp b/tests/linalg_tests.cpp index 7c81062a38..9c2494df4c 100644 --- a/tests/linalg_tests.cpp +++ b/tests/linalg_tests.cpp @@ -189,6 +189,13 @@ TEST_CASE("[mlx.core.linalg.norm] double ord") { Device::cpu) .item(), doctest::Approx(3.0)); + CHECK_EQ( + norm(x, -1, std::vector{-1, -2}, false, Device::cpu).item(), + doctest::Approx(3.0)); + CHECK_EQ( + norm(x, 1, std::vector{-1, -2}, false, Device::cpu).item(), + doctest::Approx(21.0)); + x = reshape(arange(18, float32), {2, 3, 3}); CHECK_THROWS(norm(x, 2.0, std::vector{0, 1, 2})); CHECK(allclose( @@ -283,6 +290,14 @@ TEST_CASE("[mlx.core.linalg.norm] double ord") { Device::cpu), array({3.0, 30.0})) .item()); + CHECK(allclose( + norm(x, 1, std::vector{-1, -2}, false, Device::cpu), + array({21.0, 48.0})) + .item()); + CHECK(allclose( + norm(x, -1, std::vector{-1, -2}, false, Device::cpu), + array({3.0, 30.0})) + .item()); } TEST_CASE("[mlx.core.linalg.norm] string ord") {