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

implemented isposinf and isneginf in one PR (#470)

Browse Source 
* ran precommit

* updated docs
This commit is contained in:
Yashraj Singh
2024-01-16 20:18:07 +05:30
committed by GitHub
parent a2ffea683a
commit e72458a3fa
6 changed files with 128 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files

8
mlx/ops.cpp
View File

@@ -1088,6 +1088,14 @@ array isinf(const array& a, StreamOrDevice s /* = {} */) {
return equal(a, array(std::numeric_limits<float>::infinity(), a.dtype()), s);
}
array isposinf(const array& a, StreamOrDevice s) {
return equal(a, array(std::numeric_limits<float>::infinity(), a.dtype()), s);
}
array isneginf(const array& a, StreamOrDevice s) {
return equal(a, array(-std::numeric_limits<float>::infinity(), a.dtype()), s);
}
array where(
const array& condition,
const array& x,

4
mlx/ops.h
View File

@@ -380,6 +380,10 @@ array isnan(const array& a, StreamOrDevice s = {});
array isinf(const array& a, StreamOrDevice s = {});
array isposinf(const array& a, StreamOrDevice s = {});
array isneginf(const array& a, StreamOrDevice s = {});
/** Select from x or y depending on condition. */
array where(
const array& condition,