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

Add the remainder op (#85)

Browse Source 
* Add remainder in the C++ backend
* Add the python binding and test
This commit is contained in:
Angelos Katharopoulos
2023-12-08 15:08:52 -08:00
committed by GitHub
parent 69a24e6a1e
commit 2b714714e1
14 changed files with 229 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files

8
mlx/backend/metal/kernels/binary.metal
View File

@@ -14,6 +14,13 @@ struct Divide {
template <typename T> T operator()(T x, T y) { return x / y; }
};
struct Remainder {
template <typename T> T operator()(T x, T y) { return x % y; }
template <> float operator()(float x, float y) { return fmod(x, y); }
template <> half operator()(half x, half y) { return fmod(x, y); }
template <> bfloat operator()(bfloat x, bfloat y) { return fmod(x, y); }
};
struct Equal {
template <typename T> bool operator()(T x, T y) { return x == y; }
};
@@ -363,6 +370,7 @@ instantiate_binary_types(min, Minimum)
instantiate_binary_types(mul, Multiply)
instantiate_binary_types(sub, Subtract)
instantiate_binary_types(pow, Power)
instantiate_binary_types(rem, Remainder)
// NaNEqual only needed for floating point types with boolean output
instantiate_binary_all(naneq, float16, half, bool, NaNEqual)