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

Adds round op and primitive (#203)

Browse Source
This commit is contained in:
Angelos Katharopoulos
2023-12-18 11:32:48 -08:00
committed by GitHub
parent 477397bc98
commit 4d4af12c6f
17 changed files with 187 additions and 2 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -53,6 +53,17 @@ struct SignOp {
}
};
struct RoundOp {
template <typename T>
T operator()(T x) {
return std::round(x);
}
complex64_t operator()(complex64_t x) {
return {std::round(x.real()), std::round(x.imag())};
}
};
template <typename T, typename Op>
void unary_op(const array& a, array& out, Op op) {
const T* a_ptr = a.data<T>();