From ce02855a3eb5f84f9ceff3077a728b15c49da19d Mon Sep 17 00:00:00 2001 From: Michal Klein <46717574+michalk8@users.noreply.github.com> Date: Tue, 29 Sep 2026 16:45:41 +0200 Subject: [PATCH 1/2] Fix mx.view misaligned offset --- mlx/backend/gpu/primitives.cpp | 6 ++++-- python/tests/test_ops.py | 10 ++++++++++ 2 files changed, 14 insertions(+), 2 deletions(-) diff --git a/mlx/backend/gpu/primitives.cpp b/mlx/backend/gpu/primitives.cpp index 268d6290bf..0bc19c8171 100644 --- a/mlx/backend/gpu/primitives.cpp +++ b/mlx/backend/gpu/primitives.cpp @@ -246,8 +246,10 @@ void View::eval_gpu(const std::vector& inputs, array& out) { // - type size is the same // - type size is smaller and the last axis is contiguous // - the entire array is row contiguous - if (ibytes == obytes || (obytes < ibytes && in.strides().back() == 1) || - in.flags().row_contiguous) { + // and the offset is a multiple of the output type size + if (in.offset() % obytes == 0 && + (ibytes == obytes || (obytes < ibytes && in.strides().back() == 1) || + in.flags().row_contiguous)) { auto strides = in.strides(); for (int i = 0; i < static_cast(strides.size()) - 1; ++i) { strides[i] *= ibytes; diff --git a/python/tests/test_ops.py b/python/tests/test_ops.py index a89bf7f8b0..b24e098549 100644 --- a/python/tests/test_ops.py +++ b/python/tests/test_ops.py @@ -3817,6 +3817,16 @@ def test_view(self): a_out = out.view(mx.int32) self.assertTrue(mx.array_equal(a_out, a, equal_nan=True)) + # Byte offsets that are not a multiple of the new item size + device = mx.gpu if mx.is_available(mx.gpu) else mx.cpu + a = mx.arange(64, dtype=mx.uint8) + a_np = np.arange(64, dtype=np.uint8) + with mx.stream(device): + for start in range(8): + out = a[start : start + 16].view(mx.int32) + expected = a_np[start : start + 16].view(np.int32) + self.assertTrue(np.array_equal(out + 0, expected)) # read in a kernel + def _hadamard(self, N): # Matches scipy.linalg.hadamard H = np.array([[1]], dtype=np.int64) From 1e2c3fcede4b10cda9029a0b06731ea847cc53a1 Mon Sep 17 00:00:00 2001 From: michalk8 <46717574+michalk8@users.noreply.github.com> Date: Wed, 30 Sep 2026 08:36:07 +0000 Subject: [PATCH 2/2] Update python/tests/test_ops.py Co-authored-by: Cheng --- python/tests/test_ops.py | 10 ++++------ 1 file changed, 4 insertions(+), 6 deletions(-) diff --git a/python/tests/test_ops.py b/python/tests/test_ops.py index b24e098549..a3123a6c7d 100644 --- a/python/tests/test_ops.py +++ b/python/tests/test_ops.py @@ -3818,14 +3818,12 @@ def test_view(self): self.assertTrue(mx.array_equal(a_out, a, equal_nan=True)) # Byte offsets that are not a multiple of the new item size - device = mx.gpu if mx.is_available(mx.gpu) else mx.cpu a = mx.arange(64, dtype=mx.uint8) a_np = np.arange(64, dtype=np.uint8) - with mx.stream(device): - for start in range(8): - out = a[start : start + 16].view(mx.int32) - expected = a_np[start : start + 16].view(np.int32) - self.assertTrue(np.array_equal(out + 0, expected)) # read in a kernel + for start in range(8): + out = a[start : start + 16].view(mx.int32) + expected = a_np[start : start + 16].view(np.int32) + self.assertTrue(np.array_equal(out + 0, expected)) # read in a kernel def _hadamard(self, N): # Matches scipy.linalg.hadamard