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 1 Packages Projects Releases Wiki Activity

Fix CPU sign for unsigned ints (#2024)

Browse Source 
Co-authored-by: Angelos Katharopoulos <a_katharopoulos@apple.com>
This commit is contained in:
Jesper Stemann Andersen
2025-03-31 02:56:59 +02:00
committed by GitHub
parent 28f39e9038
commit 5f5770e3a2
2 changed files with 45 additions and 3 deletions
Download Patch File Download Diff File Expand all files Collapse all files

7
mlx/backend/cpu/unary_ops.h
View File

@@ -86,13 +86,14 @@ struct Sign {
template <int N, typename T>
Simd<T, N> operator()(Simd<T, N> x) {
auto z = Simd<T, N>{0};
auto o = Simd<T, N>{1};
auto m = Simd<T, N>{-1};
if constexpr (std::is_unsigned_v<T>) {
return x != z;
return simd::select(x == z, z, o);
} else if constexpr (std::is_same_v<T, complex64_t>) {
return simd::select(x == z, x, Simd<T, N>(x / simd::abs(x)));
} else {
return simd::select(
x < z, Simd<T, N>{-1}, simd::select(x > z, Simd<T, N>{1}, z));
return simd::select(x < z, m, simd::select(x > z, o, z));
}
}
SINGLE()