comment

ml-explore · Jan 31, 2025 · eb9a9d5 · eb9a9d5
1 parent aadf66e
commit eb9a9d5
Showing 1 changed file with 4 additions and 2 deletions.
diff --git a/mlx/backend/common/indexing.cpp b/mlx/backend/common/indexing.cpp
@@ -461,7 +461,8 @@ void Scatter::eval_cpu(const std::vector<array>& inputs, array& out) {
   auto& updates = inputs.back();
 
   // Copy src into out (copy allocates memory for out)
-  copy(src, out, CopyType::General);
+  auto ctype = src.flags().row_contigous ? CopyType::Vector : CopyType::General;
+  copy(src, out, ctype);
 
   switch (src.dtype()) {
     case bool_:
@@ -621,7 +622,8 @@ void ScatterAxis::eval_cpu(const std::vector<array>& inputs, array& out) {
   auto& updates = inputs[2];
 
   // Copy src into out (copy allocates memory for out)
-  copy(src, out, CopyType::General);
+  auto ctype = src.flags().row_contigous ? CopyType::Vector : CopyType::General;
+  copy(src, out, ctype);
 
   switch (src.dtype()) {
     case bool_: