zhangyiss/mlx
1
0
Fork 0
mirror of https://github.com/ml-explore/mlx.git synced 2025-09-19 02:38:09 +08:00
Code Issues Actions Packages Projects Releases Wiki Activity

Fix round to round half-cases to even (#482)

Browse Source
This commit is contained in:
Angelos Katharopoulos
2024-01-17 15:27:23 -08:00
committed by GitHub
parent 135fd796d2
commit 90c234b7ac
4 changed files with 19 additions and 9 deletions
Download Patch File Download Diff File Expand all files Collapse all files

4
mlx/backend/common/unary.h
View File

@@ -56,11 +56,11 @@ struct SignOp {
struct RoundOp {
template <typename T>
T operator()(T x) {
return std::round(x);
return std::rint(x);
}
complex64_t operator()(complex64_t x) {
return {std::round(x.real()), std::round(x.imag())};
return {std::rint(x.real()), std::rint(x.imag())};
}
};

4
mlx/backend/metal/kernels/unary.metal
View File

@@ -134,8 +134,8 @@ struct Negative {
};
struct Round {
template <typename T> T operator()(T x) { return metal::round(x); };
template <> complex64_t operator()(complex64_t x) { return {metal::round(x.real), metal::round(x.imag)}; };
template <typename T> T operator()(T x) { return metal::rint(x); };
template <> complex64_t operator()(complex64_t x) { return {metal::rint(x.real), metal::rint(x.imag)}; };
};
struct Sigmoid {