From 9795e0ae3638b7a1d2eeeeb073fb7c770b98c901 Mon Sep 17 00:00:00 2001 From: Awni Hannun Date: Tue, 15 Jul 2025 22:10:23 -0700 Subject: [PATCH] fix cpu sort --- mlx/backend/cpu/sort.cpp | 8 ++++++-- 1 file changed, 6 insertions(+), 2 deletions(-) diff --git a/mlx/backend/cpu/sort.cpp b/mlx/backend/cpu/sort.cpp index f2243f60f..089f7c425 100644 --- a/mlx/backend/cpu/sort.cpp +++ b/mlx/backend/cpu/sort.cpp @@ -334,7 +334,9 @@ void Sort::eval_cpu(const std::vector& inputs, array& out) { auto& in = inputs[0]; // Copy input to output - CopyType ctype = in.flags().contiguous ? CopyType::Vector : CopyType::General; + CopyType ctype = (in.flags().contiguous && in.strides()[axis_] != 0) + ? CopyType::Vector + : CopyType::General; copy_cpu(in, out, ctype, stream()); auto& encoder = cpu::get_command_encoder(stream()); @@ -426,7 +428,9 @@ void Partition::eval_cpu(const std::vector& inputs, array& out) { auto& in = inputs[0]; // Copy input to output - CopyType ctype = in.flags().contiguous ? CopyType::Vector : CopyType::General; + CopyType ctype = (in.flags().contiguous && in.strides()[axis_] != 0) + ? CopyType::Vector + : CopyType::General; copy_cpu(in, out, ctype, stream()); auto& encoder = cpu::get_command_encoder(stream());