From 1c74910942f99f754b306454dd7ae545b46af6f8 Mon Sep 17 00:00:00 2001 From: Michal Klein <46717574+michalk8@users.noreply.github.com> Date: Tue, 29 Sep 2026 20:43:01 +0200 Subject: [PATCH] Fix view last axis stride --- mlx/backend/cpu/primitives.cpp | 4 ++++ mlx/backend/gpu/primitives.cpp | 4 ++++ python/tests/test_ops.py | 5 +++++ 3 files changed, 13 insertions(+) diff --git a/mlx/backend/cpu/primitives.cpp b/mlx/backend/cpu/primitives.cpp index 3299b74717..f70a2740df 100644 --- a/mlx/backend/cpu/primitives.cpp +++ b/mlx/backend/cpu/primitives.cpp @@ -380,6 +380,10 @@ void View::eval_cpu(const std::vector& inputs, array& out) { strides[i] *= ibytes; strides[i] /= obytes; } + // Last axis can have any stride + if (obytes < ibytes) { + strides.back() = 1; + } out.copy_shared_buffer( in, strides, in.flags(), in.data_size() * ibytes / obytes); } else { diff --git a/mlx/backend/gpu/primitives.cpp b/mlx/backend/gpu/primitives.cpp index 268d6290bf..d51239a2f5 100644 --- a/mlx/backend/gpu/primitives.cpp +++ b/mlx/backend/gpu/primitives.cpp @@ -253,6 +253,10 @@ void View::eval_gpu(const std::vector& inputs, array& out) { strides[i] *= ibytes; strides[i] /= obytes; } + // Last axis can have any stride + if (obytes < ibytes) { + strides.back() = 1; + } out.copy_shared_buffer( in, strides, in.flags(), in.data_size() * ibytes / obytes); } else { diff --git a/python/tests/test_ops.py b/python/tests/test_ops.py index a89bf7f8b0..fb91b37a0b 100644 --- a/python/tests/test_ops.py +++ b/python/tests/test_ops.py @@ -3817,6 +3817,11 @@ def test_view(self): a_out = out.view(mx.int32) self.assertTrue(mx.array_equal(a_out, a, equal_nan=True)) + # Size one last axis + a = mx.arange(4, dtype=mx.int32).reshape(1, 4).T + expected = np.arange(4, dtype=np.int32).reshape(4, 1).view(np.int16) + self.assertTrue(np.array_equal(a.view(mx.int16), expected)) + def _hadamard(self, N): # Matches scipy.linalg.hadamard H = np.array([[1]], dtype=np.int64)