zhangyiss/mlx
1
0
Fork 0
mirror of https://github.com/ml-explore/mlx.git synced 2025-12-16 01:49:05 +08:00
Code Issues Actions Packages Projects Releases Wiki Activity

Fix scatter_axis on single axis arrays

Browse Source
This commit is contained in:
Angelos Katharopoulos
2025-07-03 14:50:23 -07:00
parent 06a2e74eb2
commit f2994d5b29
1 changed files with 12 additions and 3 deletions
Download Patch File Download Diff File Expand all files Collapse all files

15
mlx/backend/metal/indexing.cpp
View File

@@ -1,5 +1,6 @@
// Copyright © 2023-2024 Apple Inc.
#include <fmt/format.h>
#include <iostream>
#include "mlx/backend/common/compiled.h"
#include "mlx/backend/common/utils.h"
@@ -575,9 +576,17 @@ void ScatterAxis::eval_gpu(const std::vector<array>& inputs, array& out) {
compute_encoder.set_output_array(out, 2);
// Set source info
compute_encoder.set_vector_bytes(remove_index(idx.shape(), axis_), 3);
compute_encoder.set_vector_bytes(remove_index(upd.strides(), axis_), 4);
compute_encoder.set_vector_bytes(remove_index(idx.strides(), axis_), 5);
if (ndim > 1) {
compute_encoder.set_vector_bytes(remove_index(idx.shape(), axis_), 3);
compute_encoder.set_vector_bytes(remove_index(upd.strides(), axis_), 4);
compute_encoder.set_vector_bytes(remove_index(idx.strides(), axis_), 5);
} else {
// The following will be ignored in the kernel but we still have to set
// some value so that metal validation passes.
compute_encoder.set_vector_bytes(idx.shape(), 3);
compute_encoder.set_vector_bytes(upd.strides(), 4);
compute_encoder.set_vector_bytes(idx.strides(), 5);
}
compute_encoder.set_bytes(ndim - 1, 6);
compute_encoder.set_bytes(axis_, 7);
compute_encoder.set_bytes(out.shape(axis_), 8);