zhangyiss/mlx
1
0
Fork 0
mirror of https://github.com/ml-explore/mlx.git synced 2025-07-23 18:11:17 +08:00
Code Issues Actions Packages Projects Releases Wiki Activity

Fix small sort with metal validation (#1695)

Browse Source
This commit is contained in:
Awni Hannun 2024-12-12 09:21:45 -08:00 committed by GitHub
parent 6bd28d246e
commit 9111999af3
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
2 changed files with 17 additions and 3 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

14
mlx/backend/metal/sort.cpp
View File

@ -82,9 +82,17 @@ void single_block_sort(
compute_encoder.set_bytes(out_stride_segment_axis, 6);
} else {
compute_encoder.set_bytes(nc_dim, 5);
compute_encoder.set_vector_bytes(nc_shape, 6);
compute_encoder.set_vector_bytes(in_nc_str, 7);
compute_encoder.set_vector_bytes(out_nc_str, 8);
if (nc_shape.empty()) {
int shape = 0;
int64_t stride = 0;
compute_encoder.set_bytes(shape, 6);
compute_encoder.set_bytes(stride, 7);
compute_encoder.set_bytes(stride, 8);
} else {
compute_encoder.set_vector_bytes(nc_shape, 6);
compute_encoder.set_vector_bytes(in_nc_str, 7);
compute_encoder.set_vector_bytes(out_nc_str, 8);
}
}
MTL::Size group_dims = MTL::Size(bn, 1, 1);

6
python/tests/test_ops.py
View File

@ -1956,6 +1956,12 @@ class TestOps(mlx_tests.MLXTestCase):
b_mx = mx.sort(a_mx)
self.assertTrue(np.array_equal(b_np, b_mx))
# 1D strided sort
a = mx.array([[4, 3], [2, 1], [5, 4], [3, 2]])
out = mx.argsort(a[:, 1])
expected = mx.array([1, 3, 0, 2], dtype=mx.uint32)
self.assertTrue(mx.array_equal(out, expected))
def test_partition(self):
shape = (3, 4, 5)
for dtype in ("int32", "float32"):