zhangyiss/mlx
1
0
Fork 0
mirror of https://github.com/ml-explore/mlx.git synced 2025-06-25 01:41:17 +08:00
Code Issues Actions 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
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
2 changed files with 45 additions and 3 deletions
Show Stats 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()

41
tests/ops_tests.cpp
View File

@ -1241,6 +1241,47 @@ TEST_CASE("test arithmetic unary ops") {
// bool
x = array({false, true});
CHECK(array_equal(sign(x), x).item<bool>());
// uint64
array x_uint64(
{uint64_t(0xa11cc311cb6acd70),
uint64_t(0x7a375ac3ebb533f3),
uint64_t(0x734969adf9d7190c),
uint64_t(0xb400515a4f673424)});
array expected(
{uint64_t(0x0000000000000001),
uint64_t(0x0000000000000001),
uint64_t(0x0000000000000001),
uint64_t(0x0000000000000001)});
CHECK(array_equal(sign(x_uint64), expected).item<bool>());
x_uint64 = array(
{uint64_t(0xa11cc311cb6acd70),
uint64_t(0x7a375ac3ebb533f3),
uint64_t(0x734969adf9d7190c)});
expected = array(
{uint64_t(0x0000000000000001),
uint64_t(0x0000000000000001),
uint64_t(0x0000000000000001)});
CHECK(array_equal(sign(x_uint64), expected).item<bool>());
x_uint64 =
array({uint64_t(0xa11cc311cb6acd70), uint64_t(0x7a375ac3ebb533f3)});
expected =
array({uint64_t(0x0000000000000001), uint64_t(0x0000000000000001)});
CHECK(array_equal(sign(x_uint64), expected).item<bool>());
x_uint64 = array({uint64_t(0xa11cc311cb6acd70)});
expected = array({uint64_t(0x0000000000000001)});
CHECK(array_equal(sign(x_uint64), expected).item<bool>());
x_uint64 = array({uint64_t(0xffffffffffffffff)});
expected = array({uint64_t(0x0000000000000001)});
CHECK(array_equal(sign(x_uint64), expected).item<bool>());
x_uint64 = array({uint64_t(0x0000000000000001)});
expected = array({uint64_t(0x0000000000000001)});
CHECK(array_equal(sign(x_uint64), expected).item<bool>());
}
constexpr float neginf = -std::numeric_limits<float>::infinity();